nodes/Amorano-Jovi_Capture
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

first pass window capture support

Browse Source 
tweaked region support window/monitor
This commit is contained in:
Alexander G. Morano
2025-02-22 21:15:13 -05:00
parent 09bface639
commit de8653e2ed
12 changed files with 379 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2 -1
View File
@@ -22,7 +22,8 @@ from cozy_comfyui.node import loader
PACKAGE = "JOV_CAPTURE"
WEB_DIRECTORY = "./web"
NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = loader(Path(__file__).resolve().parent,
ROOT = Path(__file__).resolve().parent
NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = loader(ROOT,
PACKAGE,
"core",
f"{PACKAGE} 📸")
core/__init__.py
+2 -6
View File
@@ -7,7 +7,7 @@ __version__ = "1.0.0"
from typing import Dict
import torch
import numpy as np
from cozy_comfyui.node import CozyImageNode
from cozy_comfyui import \
@@ -35,11 +35,7 @@ class StreamNodeHeader(CozyImageNode):
def __init__(self, *arg, **kw) -> None:
super().__init__(*arg, **kw)
self.empty = [
torch.zeros((1, MIN_IMAGE_SIZE, MIN_IMAGE_SIZE, 4), dtype=torch.uint8, device="cpu"),
torch.zeros((1, MIN_IMAGE_SIZE, MIN_IMAGE_SIZE, 3), dtype=torch.uint8, device="cpu"),
torch.zeros((1, MIN_IMAGE_SIZE, MIN_IMAGE_SIZE, 1), dtype=torch.uint8, device="cpu")
]
self.empty = np.zeros((MIN_IMAGE_SIZE, MIN_IMAGE_SIZE, 4), dtype=np.uint8)
class VideoStreamNodeHeader(StreamNodeHeader):
@classmethod
core/node_monitor.py
+29 -29
View File
@@ -1,7 +1,4 @@
"""
Jovi_Capture - http://www.github.com/amorano/Jovi_Capture
Monitor -- Capture Monitor
"""
"""Capture Monitors"""
import time
from typing import Dict
@@ -87,46 +84,49 @@ Capture frames from a desktop monitor. Supports batch processing, allowing multi
def run(self, **kw) -> RGBAMaskType:
if JOV_DOCKERENV:
return self.empty
img = cv_to_tensor_full(self.empty)
return [torch.stack(i) for i in zip(*img)]
# only allow monitor to capture single one per "batch"
monitor = parse_param(kw, "MONITOR", EnumConvertType.STRING, "NONE")[0]
try:
monitor = int(monitor.split('-')[0].strip())
except Exception:
logger.warning(f"bad monitor {monitor}")
return self.empty
images = []
batch_size = parse_param(kw, "BATCH", EnumConvertType.INT, 1, 1)[0]
# allow these to "flex" length so as to animate
monitor = parse_param(kw, "MONITOR", EnumConvertType.STRING, "NONE")
fps = parse_param(kw, "FPS", EnumConvertType.INT, 30)
xy = parse_param(kw, "XY", EnumConvertType.VEC2INT, [(0,0)], 0)
wh = parse_param(kw, "WH", EnumConvertType.VEC2INT, [(0,0)], 0)
pbar = ProgressBar(batch_size)
batch_size = [batch_size] * batch_size
params = list(zip_longest_fill(fps, xy, wh, batch_size))
size = [batch_size] * batch_size
params = list(zip_longest_fill(monitor, fps, xy, wh, size))
with mss.mss() as screen:
for idx, (fps, xy, wh, batch_size) in enumerate(params):
rate = 1. / fps
capture = screen.monitors[monitor]
width = capture['width']
height = capture['height']
width = width if wh[0] == 0 else np.clip(wh[0], 1, width)
height = height if wh[1] == 0 else np.clip(wh[1], 1, height)
region = {
'top': capture['top'] + xy[1],
'left': capture['left'] + xy[0],
'width': width,
'height': height
}
img = screen.grab(region)
img = cv2.cvtColor(np.array(img, dtype=np.uint8), cv2.COLOR_RGB2BGR)
for idx, (monitor, fps, xy, wh, size) in enumerate(params):
try:
monitor = int(monitor.split('-')[0].strip())
except Exception:
logger.warning(f"bad monitor {monitor}")
img = self.empty
else:
capture = screen.monitors[monitor]
width = capture['width']
height = capture['height']
width = width if wh[0] == 0 else np.clip(wh[0], 1, width)
height = height if wh[1] == 0 else np.clip(wh[1], 1, height)
region = {
'top': capture['top'] + xy[1],
'left': capture['left'] + xy[0],
'width': width,
'height': height
}
img = screen.grab(region)
img = cv2.cvtColor(np.array(img, dtype=np.uint8), cv2.COLOR_RGB2BGR)
images.append(cv_to_tensor_full(img))
pbar.update_absolute(idx)
if batch_size > 1:
rate = 1. / fps
time.sleep(rate)
return [torch.stack(i) for i in zip(*images)]
core/node_remote.py
+1 -1
View File
@@ -25,7 +25,7 @@ from . import StreamNodeHeader
# ==============================================================================
class RemoteSteamReader(StreamNodeHeader):
NAME = "REMOTE URL"
NAME = "REMOTE"
DESCRIPTION = """
Capture frames from a URL. Supports batch processing, allowing multiple frames to be captured simultaneously. The node provides options for configuring the source, resolution, frame rate, zoom, orientation, and interpolation method. Additionally, it supports capturing frames from multiple monitors or windows simultaneously.
"""
core/node_webcam.py
+16 -21
View File
@@ -1,12 +1,8 @@
"""
Jovi_Capture - http://www.github.com/amorano/Jovi_Capture
Capture -- WEBCAM, REMOTE URLS
"""
"""Capture -- WEBCAM"""
import os
import time
from typing import Any, Dict, List, Tuple
from typing import Any, Dict, List
import cv2
import torch
@@ -21,6 +17,7 @@ from cozy_comfyui import \
EnumConvertType, \
deep_merge, parse_param
from cozy_comfyui import RGBAMaskType
from cozy_comfyui.image.convert import cv_to_tensor_full
from . import VideoStreamNodeHeader
@@ -37,25 +34,25 @@ JOV_SCAN_DEVICES = os.getenv("JOV_SCAN_DEVICES", "False").lower() in ['1', 'true
# === SUPPORT ===
# ==============================================================================
def cameraList() -> List[str]:
def camera_list() -> List[str]:
idx = 0
failed = 0
cameraList = []
camera_list = []
while failed < 2:
cap = cv2.VideoCapture(idx)
if cap.isOpened():
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
f = int(cap.get(cv2.CAP_PROP_FPS))
cameraList.append(f"{idx} - {w}x{h}x{f}")
camera_list.append(f"{idx} - {w}x{h}x{f}")
cap.release()
else:
failed += 1
idx += 1
if len(cameraList) == 0:
cameraList = ["NONE"]
return cameraList
if len(camera_list) == 0:
camera_list = ["NONE"]
return camera_list
# ==============================================================================
# === API ROUTE ===
@@ -64,7 +61,7 @@ def cameraList() -> List[str]:
@PromptServer.instance.routes.get(f"/{PACKAGE.lower()}/camera")
async def route_cameraList(req) -> Any:
# load the camera list here..
CameraStreamReader.CAMERAS = cameraList()
CameraStreamReader.CAMERAS = camera_list()
return web.json_response(CameraStreamReader.CAMERAS)
# ==============================================================================
@@ -149,7 +146,7 @@ Capture frames from a web camera. Supports batch processing, allowing multiple f
d = super().INPUT_TYPES()
if cls.CAMERAS is None:
cls.CAMERAS = cameraList() if JOV_SCAN_DEVICES else ["NONE"]
cls.CAMERAS = camera_list() if JOV_SCAN_DEVICES else ["NONE"]
return deep_merge({
"optional": {
@@ -161,14 +158,15 @@ Capture frames from a web camera. Supports batch processing, allowing multiple f
}
}, d)
def run(self, **kw) -> Tuple[torch.Tensor, ...]:
def run(self, **kw) -> RGBAMaskType:
# need to see if we have a device...
url = parse_param(kw, "CAMERA", EnumConvertType.STRING, "")[0]
try:
url = int(url.split('-')[0].strip())
except Exception:
logger.warning(f"bad camera url {url}")
return self.empty
img = cv_to_tensor_full(self.empty)
return [torch.stack(i) for i in zip(*img)]
if self.device is None:
self.device = MediaStreamCamera()
@@ -176,10 +174,6 @@ Capture frames from a web camera. Supports batch processing, allowing multiple f
self.device.timeout = parse_param(kw, "TIMEOUT", EnumConvertType.INT, 5, 1, 30)[0]
self.device.url = url
#wh = parse_param(kw, "WH", EnumConvertType.VEC2INT, [640, 480], 160)[0]
#self.device.width = wh[0]
#self.device.height = wh[1]
images = []
self.device.fps = parse_param(kw, "FPS", EnumConvertType.INT, 30)[0]
batch_size = parse_param(kw, "BATCH", EnumConvertType.INT, 1, 1)[0]
@@ -202,7 +196,8 @@ Capture frames from a web camera. Supports batch processing, allowing multiple f
break
if time.perf_counter() - start_time > self.device.timeout:
logger.error("could not capture device")
return self.empty
img = self.empty
break
images.append(cv_to_tensor_full(img))
if batch_size > 1:
core/node_window.py
+252 -57
View File
@@ -1,25 +1,237 @@
"""
Jovi_Capture - http://www.github.com/amorano/Jovi_Capture
Window -- Stream dekstop window
"""
"""Capture Dekstop Window"""
import re
import json
import time
from typing import Dict, Tuple
import platform
from typing import Any, Dict, Optional, Tuple
import cv2
import torch
import numpy as np
import pywinctl as pwc
from aiohttp import web
from loguru import logger
from comfy.utils import ProgressBar
from server import PromptServer
from cozy_comfyui import \
EnumConvertType, \
deep_merge, parse_param
deep_merge, parse_param, zip_longest_fill
from cozy_comfyui import RGBAMaskType
from cozy_comfyui.image import ImageType
from cozy_comfyui.image.convert import cv_to_tensor_full
if platform.system() == "Windows":
import win32gui
import win32ui
from ctypes import windll
elif platform.system() == "Darwin":
from Quartz import *
elif platform.system() == "Linux":
from Xlib import display, X
from Xlib.ext import composite
from . import StreamNodeHeader
from .. import ROOT, PACKAGE
# ==============================================================================
# === INITIALIZE ===
# ==============================================================================
try:
with open(f"{ROOT}/skip.json", "r") as fp:
IGNORE_LIST = json.load(fp)
IGNORE_LIST['regex'] = [re.compile(x) for x in IGNORE_LIST['regex']]
except Exception as e:
logger.error(e)
IGNORE_LIST = {
"full": [],
"regex": []
}
# ==============================================================================
# === SUPPORT ===
# ==============================================================================
def window_list() -> Dict[int, str]:
"""List all draggable, user-usable windows with their handles and titles."""
windows = pwc.getAllWindows()
valid_windows = {}
for win in windows:
if win.isVisible and not win.isMinimized and \
win.width>0 and win.height>0 \
and win.title not in IGNORE_LIST['full'] and \
not any(pattern.search(win.title) for pattern in IGNORE_LIST['regex']):
valid_windows[win.title] = win.getHandle()
return valid_windows
def window_capture(hwnd: int, client_area_only: bool=False, region: Optional[Tuple[int,...]]=None) -> ImageType:
"""
Capture a window or region within a window.
Args:
hwnd: Window handle
client_area_only: If True, captures only the client area without borders/decorations
region: Optional (x, y, width, height) tuple specifying region within window to capture
Returns:
ImageType: Captured image in RGBA format
"""
system = platform.system()
if system == "Windows":
# Get correct window rect based on capture mode
if client_area_only:
rect = win32gui.GetClientRect(hwnd)
left, top = win32gui.ClientToScreen(hwnd, (0, 0))
right = left + rect[2]
bottom = top + rect[3]
else:
left, top, right, bottom = win32gui.GetWindowRect(hwnd)
width = right - left
height = bottom - top
# Adjust for region if specified
if region:
rx, ry, rw, rh = region
left += rx
top += ry
width = min(rw, width - rx)
height = min(rh, height - ry)
window_dc = win32gui.GetWindowDC(hwnd)
dc = win32ui.CreateDCFromHandle(window_dc)
compatible_dc = dc.CreateCompatibleDC()
try:
bitmap = win32ui.CreateBitmap()
bitmap.CreateCompatibleBitmap(dc, width, height)
compatible_dc.SelectObject(bitmap)
# Set the correct source coordinates for BitBlt
if client_area_only or region:
compatible_dc.BitBlt((0, 0), (width, height), dc, (rx if region else 0, ry if region else 0), win32con.SRCCOPY)
else:
windll.user32.PrintWindow(hwnd, compatible_dc.GetSafeHdc(), 2)
bmpstr = bitmap.GetBitmapBits(True)
img = np.frombuffer(bmpstr, dtype='uint8')
img = img.reshape((height, width, 4))
finally:
dc.DeleteDC()
compatible_dc.DeleteDC()
win32gui.ReleaseDC(hwnd, window_dc)
win32gui.DeleteObject(bitmap.GetHandle())
return cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
elif system == "Darwin":
# Get window info
window_list = CGWindowListCopyWindowInfo(
kCGWindowListOptionIncludingWindow,
hwnd
)
window_info = window_list[0]
# Get bounds
bounds = window_info[kCGWindowBounds]
if client_area_only:
# Adjust bounds to exclude title bar and borders
# Note: This is an approximation, as macOS doesn't have a direct equivalent
bounds.origin.y += 22 # Typical title bar height
bounds.size.height -= 22
if region:
rx, ry, rw, rh = region
bounds.origin.x += rx
bounds.origin.y += ry
bounds.size.width = min(rw, bounds.size.width - rx)
bounds.size.height = min(rh, bounds.size.height - ry)
# Create image of window contents
image = CGWindowListCreateImage(
bounds,
kCGWindowListOptionIncludingWindow,
hwnd,
kCGWindowImageBoundsIgnoreFraming | kCGWindowImageShouldBeOpaque
)
dataProvider = CGImageGetDataProvider(image)
data = dataProvider.copy()
width = CGImageGetWidth(image)
height = CGImageGetHeight(image)
img = np.frombuffer(data, dtype=np.uint8)
return img.reshape((height, width, 4))
elif system == "Linux":
d = display.Display()
window = d.create_resource_object('window', hwnd)
if client_area_only:
# Get window properties to find decorations
prop = window.get_full_property(
d.intern_atom('_NET_FRAME_EXTENTS'),
X.AnyPropertyType
)
if prop:
# left, right, top, bottom
frame_extents = prop.value
x = frame_extents[0]
y = frame_extents[2]
geom = window.get_geometry()
width = geom.width - (frame_extents[0] + frame_extents[1])
height = geom.height - (frame_extents[2] + frame_extents[3])
else:
geom = window.get_geometry()
x, y = 0, 0
width, height = geom.width, geom.height
else:
geom = window.get_geometry()
x, y = 0, 0
width, height = geom.width, geom.height
if region:
rx, ry, rw, rh = region
x += rx
y += ry
width = min(rw, width - rx)
height = min(rh, height - ry)
try:
composite.composite_redirect_window(d, window, True)
pixmap = window.create_pixmap(width, height, window.get_attributes().depth)
gc = pixmap.create_gc()
# Copy the specified region
window.composite_name_window_pixmap()
gc.copy_area(window, pixmap, x, y, 0, 0, width, height)
image = pixmap.get_image(0, 0, width, height, X.ZPixmap, 0xffffffff)
img = np.frombuffer(image.data, dtype=np.uint8)
img = img.reshape((height, width, 4))
finally:
gc.free()
pixmap.free()
return cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
# ==============================================================================
# === API ROUTE ===
# ==============================================================================
@PromptServer.instance.routes.get(f"/{PACKAGE.lower()}/window")
async def route_windowList(req) -> Any:
WindowStreamReader.WINDOWS = window_list()
return web.json_response(WindowStreamReader.WINDOWS)
# ==============================================================================
# === NODE ===
@@ -30,68 +242,51 @@ class WindowStreamReader(StreamNodeHeader):
DESCRIPTION = """
Capture frames from a dekstop window. Supports batch processing, allowing multiple frames to be captured simultaneously. The node provides options for configuring the source, resolution, frame rate, zoom, orientation, and interpolation method. Additionally, it supports capturing frames from multiple monitors or windows simultaneously.
"""
WINDOWS = None
@classmethod
def INPUT_TYPES(cls) -> Dict[str, str]:
d = super().INPUT_TYPES()
if cls.WINDOWS is None:
cls.WINDOWS = window_list()
default = ""
keys = []
if len(cls.WINDOWS):
keys = list(cls.WINDOWS.keys())
default = keys[0]
return deep_merge({
"optional": {
"WINDOW": (keys, {"default": default, "tooltip": "Window to capture"}),
"XY": ("VEC2INT", {"default": (0, 0), "mij": 0, "label": ["TOP", "LEFT"], "tooltip": "Top, Left position"}),
"WH": ("VEC2INT", {"default": (0, 0), "mij": 0, "label": ["WIDTH", "HEIGHT"], "tooltip": "Width and Height"}),
"CLIENT": ("BOOLEAN", {"default": False, "tooltip": "Only capture the client area -- no scrollbars or menus"}),
}
}, d)
def __init__(self, *arg, **kw) -> None:
super().__init__(*arg, **kw)
self.__device = None
def run(self, **kw) -> Tuple[torch.Tensor, torch.Tensor]:
wait = parse_param(kw, "WAIT", EnumConvertType.BOOLEAN, False)[0]
if wait:
return self.__last
def run(self, **kw) -> RGBAMaskType:
images = []
batch_size, rate = parse_param(kw, "BATCH", EnumConvertType.VEC2INT, [(1, 30)], 1)[0]
batch_size = parse_param(kw, "BATCH", EnumConvertType.INT, 1, 1)[0]
window = parse_param(kw, "WINDOW", EnumConvertType.STRING, "")
fps = parse_param(kw, "FPS", EnumConvertType.INT, 30)
xy = parse_param(kw, "XY", EnumConvertType.VEC2INT, [(0,0)], 0)
wh = parse_param(kw, "WH", EnumConvertType.VEC2INT, [(0,0)], 0)
client = parse_param(kw, "CLIENT", EnumConvertType.BOOLEAN, False)
pbar = ProgressBar(batch_size)
rate = 1. / rate
size = [batch_size] * batch_size
params = list(zip_longest_fill(window, fps, xy, wh, client, size))
for idx, (window, fps, xy, wh, client, size) in enumerate(params):
window = self.WINDOWS[window]
region = None
if (img := window_capture(window, client, region)) is None:
img = self.empty
camera = parse_param(kw, "CAMERA", EnumConvertType.STRING, "")[0]
camera = camera.split('-')[0].strip()
try:
_ = int(camera)
camera = str(camera)
except:
camera = ""
images.append(cv_to_tensor_full(img))
if batch_size > 1:
rate = 1. / fps
time.sleep(rate)
pbar.update_absolute(idx)
# timeout and try again?
if self.__capturing > 0 and time.perf_counter() - self.__capturing > 3000:
logger.error(f'timed out {self.__url}')
self.__capturing = 0
self.__url = ""
if self.__device is not None:
self.__capturing = 0
if wait:
self.__device.pause()
else:
self.__device.play()
fps = parse_param(kw, "FPS", EnumConvertType.INT, 30)[0]
self.__device.fps = fps
self.__device.zoom = parse_param(kw, "ZOOM", EnumConvertType.FLOAT, 0, 0, 1)[0]
for idx in range(batch_size):
img = self.__device.frame
if img is None:
images.append(self.__empty)
else:
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGRA)
images.append(cv_to_tensor_full(img))
pbar.update_absolute(idx)
if batch_size > 1:
time.sleep(rate)
if len(images) == 0:
images.append(self.__empty)
self.__last = [torch.stack(i) for i in zip(*images)]
return self.__last
return [torch.stack(i) for i in zip(*images)]
node_list.json
+6
View File
@@ -0,0 +1,6 @@
{
"CAMERA (JOV_CAPTURE)": "Capture frames from a web camera",
"MONITOR (JOV_CAPTURE)": "Capture frames from a desktop monitor",
"REMOTE (JOV_CAPTURE)": "Capture frames from a URL",
"WINDOW (JOV_CAPTURE)": "Capture frames from a dekstop window"
}
pyproject.toml
+4 -1
View File
@@ -22,7 +22,10 @@ dependencies = [
"numpy>=1.26.4,<2.0.0; python_version <= '3.11'",
"numpy>=2.0.0; python_version >= '3.12'",
"opencv-contrib-python",
"Pillow"
"Pillow",
"pyobjc-framework-Quartz; platform_system=='Darwin'",
"pywin32; platform_system=='Windows'",
"Xlib; platform_system=='Linux'"
]
[project.urls]
requirements.txt
+4 -1
View File
@@ -4,4 +4,7 @@ mss
numpy>=1.26.4,<2.0.0; python_version <= '3.11'
numpy>=2.0.0; python_version >= '3.12'
opencv-contrib-python
Pillow
Pillow
pyobjc-framework-Quartz; platform_system=='Darwin'
pywin32; platform_system=='Windows'
Xlib; platform_system=='Linux'
skip.json
+27
View File
@@ -0,0 +1,27 @@
{
"full": [
"",
"Address band toolbar",
"Calculator",
"Chrome Legacy Window",
"FolderView",
"Microsoft Text Input Application",
"Namespace Tree Control",
"Navigation buttons",
"Program Manager",
"Ribbon",
"Running applications",
"Settings",
"Shellview",
"Start",
"Tree View",
"UIRibbonDockTop",
"Up band toolbar",
"User Promoted Notification Area"
],
"regex": [
"System Clock,.*",
"Action Center,.*",
"Address:.*"
]
}
web/node_webcam.js
-1
View File
@@ -25,7 +25,6 @@ app.registerExtension({
var data = await api_get("/jov_capture/camera");
widget_camera.options.values = data;
widget_camera.value = data[0];
console.info(widget_camera)
app.canvas.setDirty(true);
});
return me;
web/node_window.js
+36
View File
@@ -0,0 +1,36 @@
/**
* File: node_window.js
* Project: jov_capture
*/
import { app } from "../../../scripts/app.js";
import { api_get } from './util_jov.js'
const _id = "WINDOW (JOV_CAPTURE)";
app.registerExtension({
name: 'jov_capture.node.' + _id,
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (nodeData.name !== _id) {
return
}
const onNodeCreated = nodeType.prototype.onNodeCreated
nodeType.prototype.onNodeCreated = function () {
const me = onNodeCreated?.apply(this);
const widget_window = this.widgets.find(w => w.name == 'WINDOW');
this.addWidget('button', 'REFRESH WINDOW LIST', 'refresh', async () => {
var data = await api_get("/jov_capture/window");
widget_window.options.values = Object.keys(data);
widget_window.value = widget_window.options.values[0];
console.info(widget_window)
app.canvas.setDirty(true);
});
return me;
}
return nodeType;
}
});