diff --git a/RH_Utils.py b/RH_Utils.py index 73f6648..79086a0 100644 --- a/RH_Utils.py +++ b/RH_Utils.py @@ -1,3 +1,13 @@ +import os +import numpy as np +import json +import folder_paths +import zipfile +import shutil +import numpy as np +import torch +from PIL import Image, ImageOps + class AllTrue(str): def __init__(self, representation=None) -> None: self.repr = representation @@ -55,3 +65,73 @@ class AnyToStringNode: else: # For non-string types, directly convert to string return (str(anything),) + +class RH_Extract_Image_From_List(): + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images list"}), + "image_index": ("INT", {"default": 0 }), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "rh_extract_image" + + OUTPUT_NODE = False + + CATEGORY = "RunningHub" + + def rh_extract_image(self, images, image_index): + out = images[int(image_index)].unsqueeze(0) + return (out,) + +class RH_Batch_Images_From_List(): + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images list"}), + "image_indices": ("STRING", {"default":"0-3,4,5-7","tooltip": "Some like 0-2, 3, 4-5. Leaving it empty means selecting all."}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "rh_batch_images" + + OUTPUT_NODE = False + + CATEGORY = "RunningHub" + + def rh_batch_images(self, images, image_indices): + image_indices = image_indices.replace(" ", "") + out = [] + if image_indices == "": + out = images + image_indices = image_indices.split(',') + for index in image_indices: + if '-' in index: + sindex = index.split('-') + out.extend(images[int(sindex[0]):int(sindex[1])+1]) + else: + out.append(images[int(index)]) + batchsize = len(out) + max_height = max(image.shape[0] for image in out) + max_width = max(image.shape[1] for image in out) + max_channels = max(image.shape[2] for image in out) + batch_images = torch.zeros([batchsize, max_height, max_width, max_channels]) + for (batch_number, image) in enumerate(out): + h, w, c = image.shape + batch_images[batch_number, 0:h, 0:w, 0:c] = image + return (batch_images,) \ No newline at end of file diff --git a/__init__.py b/__init__.py index ee665cf..a1c5b60 100644 --- a/__init__.py +++ b/__init__.py @@ -2,7 +2,7 @@ from .RH_SettingsNode import SettingsNode from .RH_NodeInfoListNode import NodeInfoListNode from .RH_ExecuteNode import ExecuteNode from .RH_ImageUploaderNode import ImageUploaderNode -from .RH_Utils import AnyToStringNode +from .RH_Utils import * @@ -12,6 +12,8 @@ NODE_CLASS_MAPPINGS = { "RH_ExecuteNode": ExecuteNode, "RH_ImageUploaderNode": ImageUploaderNode, "RH_Utils": AnyToStringNode, + "RH_ExtractImage": RH_Extract_Image_From_List, + "RH_BatchImages": RH_Batch_Images_From_List, } @@ -21,7 +23,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "RH_ExecuteNode": "RH Execute", "RH_ImageUploaderNode": "RH Image Uploader", "RH_Utils": "RH Anything to String", - + "RH_ExtractImage": "RH Extract Image From ImageList", + "RH_BatchImages": "RH Batch Images From ImageList", } + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS",]