xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2026-05-29 06:57:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add option for batch capture

Browse Source
This commit is contained in:
kijai 2023-11-30 01:38:41 +02:00
parent f542cdb950
commit 5fc155c073
1 changed files with 18 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
nodes.py
View File

@ -2414,7 +2414,7 @@ class SplitBboxes:
return (bboxes_a, bboxes_b,) return (bboxes_a, bboxes_b,)
from PIL import ImageGrab from PIL import ImageGrab
import time
class ImageGrabPIL: class ImageGrabPIL:
@classmethod @classmethod
@ -2435,22 +2435,27 @@ class ImageGrabPIL:
"y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), "y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}),
"width": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}), "width": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}),
"height": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}), "height": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}),
"num_frames": ("INT", {"default": 1,"min": 1, "max": 255, "step": 1}),
"delay": ("FLOAT", {"default": 0.1,"min": 0.0, "max": 10.0, "step": 0.1}),
}, },
} }
def screencap(self, x, y, width, height): def screencap(self, x, y, width, height, num_frames, delay):
captures = []
# Define the bounding box of the area you want to capture
bbox = (x, y, x + width, y + height) bbox = (x, y, x + width, y + height)
# Capture the screen for _ in range(num_frames):
screen_capture = ImageGrab.grab(bbox=bbox) # Capture screen
screen_capture = ImageGrab.grab(bbox=bbox)
screen_capture_torch = torch.tensor(np.array(screen_capture), dtype=torch.float32) / 255.0
screen_capture_torch = screen_capture_torch.unsqueeze(0)
captures.append(screen_capture_torch)
# Convert the PIL Image directly to a PyTorch tensor if that's the desired final format # Wait for a short delay if more than one frame is to be captured
screen_capture_torch = torch.tensor(np.array(screen_capture), dtype=torch.float32) / 255.0 if num_frames > 1:
screen_capture_torch = screen_capture_torch.unsqueeze(0) # Permute to have channel-first format and add batch dimension time.sleep(delay)
return (screen_capture_torch,) return (torch.cat(captures, dim=0),)
NODE_CLASS_MAPPINGS = { NODE_CLASS_MAPPINGS = {