diff --git a/nodes.py b/nodes.py index 5d104e6..230d60d 100644 --- a/nodes.py +++ b/nodes.py @@ -2131,7 +2131,6 @@ class WidgetToString: workflow = extra_pnginfo["workflow"] results = [] for node in workflow["nodes"]: - print(node) node_id = node["id"] if node_id != id: @@ -2414,6 +2413,46 @@ class SplitBboxes: return (bboxes_a, bboxes_b,) +from PIL import ImageGrab + +class ImageGrabPIL: + + @classmethod + def IS_CHANGED(cls): + + return + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "screencap" + CATEGORY = "KJNodes/experimental" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "x": ("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}), + "height": ("INT", {"default": 512,"min": 0, "max": 4096, "step": 1}), + }, + } + + def screencap(self, x, y, width, height): + + # Define the bounding box of the area you want to capture + bbox = (x, y, x + width, y + height) + + # Capture the screen + screen_capture = ImageGrab.grab(bbox=bbox) + + # Convert the PIL Image directly to a PyTorch tensor if that's the desired final format + screen_capture_torch = torch.tensor(np.array(screen_capture), dtype=torch.float32) / 255.0 + screen_capture_torch = screen_capture_torch.unsqueeze(0) # Permute to have channel-first format and add batch dimension + + return (screen_capture_torch,) + + NODE_CLASS_MAPPINGS = { "INTConstant": INTConstant, "FloatConstant": FloatConstant, @@ -2457,6 +2496,7 @@ NODE_CLASS_MAPPINGS = { "CreateMagicMask": CreateMagicMask, "BboxToInt": BboxToInt, "SplitBboxes": SplitBboxes, + "ImageGrabPIL": ImageGrabPIL } NODE_DISPLAY_NAME_MAPPINGS = { "INTConstant": "INT Constant", @@ -2500,4 +2540,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CreateMagicMask": "CreateMagicMask", "BboxToInt": "BboxToInt", "SplitBboxes": "SplitBboxes", + "ImageGrabPIL": "ImageGrabPIL" } \ No newline at end of file