diff --git a/example_workflows/Test_Load_Dataset.json b/example_workflows/Test_Load_Dataset.json new file mode 100644 index 0000000..c8a9f66 --- /dev/null +++ b/example_workflows/Test_Load_Dataset.json @@ -0,0 +1 @@ +{"id":"005f9807-4d92-44a8-8a4a-178d84b8dcb7","revision":0,"last_node_id":15,"last_link_id":11,"nodes":[{"id":3,"type":"PreviewAny","pos":[2370,1570],"size":[620,88],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"*","link":3}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.59","Node name for S&R":"PreviewAny","ue_properties":{"version":"7.0.1","widget_ue_connectable":{}}},"widgets_values":[]},{"id":10,"type":"SET_ImageLoad","pos":[2360,1800],"size":[276.734375,314],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"file_name","name":"file_name","type":"STRING","widget":{"name":"file_name"},"link":6},{"localized_name":"embed_transparency","name":"embed_transparency","shape":7,"type":"BOOLEAN","widget":{"name":"embed_transparency"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[7]},{"localized_name":"alpha_mask","name":"alpha_mask","type":"MASK","links":null}],"properties":{"cnr_id":"image-misc","ver":"1078a8a9e0702d28c73f6a3cbf0a90e5fb74b16a","Node name for S&R":"SET_ImageLoad","ue_properties":{"version":"7.0.1","widget_ue_connectable":{"file_name":true,"embed_transparency":true}}},"widgets_values":["",false]},{"id":11,"type":"PreviewImage","pos":[2770,1820],"size":[320,258],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":7}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.59","Node name for S&R":"PreviewImage","ue_properties":{"version":"7.0.1","widget_ue_connectable":{}}},"widgets_values":[]},{"id":13,"type":"PreviewImage","pos":[2770,2170],"size":[320,258],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":9}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.59","Node name for S&R":"PreviewImage","ue_properties":{"version":"7.0.1","widget_ue_connectable":{}}},"widgets_values":[]},{"id":12,"type":"SET_ImageLoad","pos":[2390,2150],"size":[276.734375,314],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"file_name","name":"file_name","type":"STRING","widget":{"name":"file_name"},"link":10},{"localized_name":"embed_transparency","name":"embed_transparency","shape":7,"type":"BOOLEAN","widget":{"name":"embed_transparency"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[9]},{"localized_name":"alpha_mask","name":"alpha_mask","type":"MASK","links":null}],"properties":{"cnr_id":"image-misc","ver":"1078a8a9e0702d28c73f6a3cbf0a90e5fb74b16a","Node name for S&R":"SET_ImageLoad","ue_properties":{"version":"7.0.1","widget_ue_connectable":{"file_name":true,"embed_transparency":true}}},"widgets_values":["",false]},{"id":4,"type":"SET_ImageDataset","pos":[1790,1580],"size":[270,218],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"source","name":"source","type":"STRING","widget":{"name":"source"},"link":null},{"localized_name":"pattern","name":"pattern","type":"STRING","widget":{"name":"pattern"},"link":null},{"localized_name":"destination","name":"destination","type":"STRING","widget":{"name":"destination"},"link":null},{"localized_name":"dest_ext","name":"dest_ext","type":"STRING","widget":{"name":"dest_ext"},"link":null},{"localized_name":"reference","name":"reference","shape":7,"type":"STRING","widget":{"name":"reference"},"link":null},{"localized_name":"sort_method","name":"sort_method","shape":7,"type":"COMBO","widget":{"name":"sort_method"},"link":null}],"outputs":[{"localized_name":"images","name":"images","shape":6,"type":"STRING","links":[3,6]},{"localized_name":"results","name":"results","shape":6,"type":"STRING","links":[]},{"localized_name":"references","name":"references","shape":6,"type":"STRING","links":[10]}],"properties":{"cnr_id":"image-misc","ver":"1d1b99ac4989ba2a5df143f6c81913ad20f3db0b","Node name for S&R":"SET_ImageDataset","ue_properties":{"version":"7.0.1","widget_ue_connectable":{"source":true,"pattern":true,"destination":true,"dest_ext":true,"reference":true,"sort_method":true}}},"widgets_values":["./DIS5K/DIS-VD/im",".*","./DIS5K/DIS-VD","png","./DIS5K/DIS-VD/gt","None"]},{"id":14,"type":"SET_ImageDownload","pos":[1640,2350],"size":[276.734375,362],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image_bypass","name":"image_bypass","shape":7,"type":"IMAGE","link":null},{"localized_name":"mask_bypass","name":"mask_bypass","shape":7,"type":"MASK","link":null},{"localized_name":"base_url","name":"base_url","type":"STRING","widget":{"name":"base_url"},"link":null},{"localized_name":"filename","name":"filename","type":"STRING","widget":{"name":"filename"},"link":null},{"localized_name":"local_name","name":"local_name","shape":7,"type":"STRING","widget":{"name":"local_name"},"link":null},{"localized_name":"embed_transparency","name":"embed_transparency","shape":7,"type":"BOOLEAN","widget":{"name":"embed_transparency"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[11]},{"localized_name":"alpha_mask","name":"alpha_mask","type":"MASK","links":null}],"properties":{"cnr_id":"image-misc","ver":"1078a8a9e0702d28c73f6a3cbf0a90e5fb74b16a","Node name for S&R":"SET_ImageDownload"},"widgets_values":["https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/","audioseparation_logo.jpg","",false]},{"id":15,"type":"PreviewImage","pos":[2348.124755859375,2502.319580078125],"size":[140,246],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":11}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.59","Node name for S&R":"PreviewImage"},"widgets_values":[]}],"links":[[3,4,0,3,0,"*"],[6,4,0,10,0,"STRING"],[7,10,0,11,0,"IMAGE"],[9,12,0,13,0,"IMAGE"],[10,4,2,12,0,"STRING"],[11,14,0,15,0,"IMAGE"]],"groups":[],"config":{},"extra":{"ue_links":[],"ds":{"scale":0.9261636231874897,"offset":[-1503.5414324598712,-1660.8292931230108]},"links_added_by_ue":[]},"version":0.4} \ No newline at end of file diff --git a/src/nodes/helpers.py b/src/nodes/helpers.py new file mode 100644 index 0000000..907e74b --- /dev/null +++ b/src/nodes/helpers.py @@ -0,0 +1,134 @@ +import numpy as np +import os +from PIL import Image, ImageOps, ImageSequence +from PIL import ImageFile, UnidentifiedImageError +import torch +from . import main_logger + +try: + # We need to import the built-in LoadImage class for ImageDownload + from nodes import LoadImage + from folder_paths import get_input_directory + has_load_image = hasattr(LoadImage, "load_image") +except Exception: + has_load_image = False + +logger = main_logger + + +def pillow(fn, arg): + prev_value = None + try: + x = fn(arg) + except (OSError, UnidentifiedImageError, ValueError): # PIL issues #4472 and #2445, also fixes ComfyUI issue #3416 + prev_value = ImageFile.LOAD_TRUNCATED_IMAGES + ImageFile.LOAD_TRUNCATED_IMAGES = True + x = fn(arg) + finally: + if prev_value is not None: + ImageFile.LOAD_TRUNCATED_IMAGES = prev_value + return x + + +class CustomLoadImage(object): + def load_image(self, image): + """ ComfyUI 0.3.59 loader """ + image_path = image + + img = pillow(Image.open, image_path) + + output_images = [] + output_masks = [] + w, h = None, None + + excluded_formats = ['MPO'] + + for i in ImageSequence.Iterator(img): + i = pillow(ImageOps.exif_transpose, i) + + if i.mode == 'I': + i = i.point(lambda i: i * (1 / 255)) + image = i.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + elif i.mode == 'P' and 'transparency' in i.info: + mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + output_images.append(image) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and img.format not in excluded_formats: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) + else: + output_image = output_images[0] + output_mask = output_masks[0] + + return (output_image, output_mask) + + +def load_image_wrapper(file_name, embed_transparency, disp_name=None): + disp_name = disp_name or file_name + + # --- REUSE ComfyUI's LoadImage LOGIC --- + try: + if has_load_image: + # Instantiate the built-in LoadImage node + loader_instance = LoadImage() + + # The LoadImage node's `load_image` method expects the filename as passed + # by the ComfyUI widget, which is just the filename. It internally + # resolves the path using folder_paths. + + logger.debug(f"Calling built-in LoadImage.load_image() with filename: '{file_name}'") + else: + # Instantiate the built-in LoadImage node + loader_instance = CustomLoadImage() + logger.debug(f"Calling our CustomLoadImage.load_image() with filename: '{file_name}'") + + # Call the method and return its result directly + result = loader_instance.load_image(file_name) + # Create an RGBA image if needed + if embed_transparency: + image, mask = result + # Expand the mask to (b, h, w, 1) + mask = mask[..., None] + # Concatenate image and mask into (b, h, w, 4) + image_with_alpha = torch.cat([image, 1.0 - mask], dim=-1) + result = (image_with_alpha, mask) + # This information is for the preview, as we are an output node and we return images + # they will be displayed in our node. Quite simple. + if os.path.isabs(file_name): + ff_name = os.path.relpath(file_name, get_input_directory()) + fname = os.path.basename(ff_name) + dname = os.path.dirname(ff_name) + else: + fname = file_name + dname = "" + downloaded_file = { + "images": [{ + "filename": fname, + "subfolder": dname, + "type": "input" # We stored the file in the "input" folder + }] + } + return {"ui": downloaded_file, "result": result} + + except Exception as e: + logger.error(f"Failed to load image '{disp_name}' using built-in LoadImage node: {e}", exc_info=True) + # Re-raise to make the error visible in ComfyUI + raise IOError(f"Could not load the image file '{disp_name}' using the standard loader. " + "It may be corrupt or in an unsupported format.") from e diff --git a/src/nodes/nodes_img.py b/src/nodes/nodes_img.py index 054db31..89e9518 100644 --- a/src/nodes/nodes_img.py +++ b/src/nodes/nodes_img.py @@ -10,26 +10,31 @@ from copy import deepcopy import numpy as np import os +from pathlib import Path from PIL import Image, ImageDraw, ImageFont # Import the Python Imaging Library +import re from seconohe.apply_mask import apply_mask from seconohe.foreground_estimation.affce import affce from seconohe.foreground_estimation.fmlfe import fmlfe, IMPL_PRIORITY from seconohe.downloader import download_file from seconohe.color import color_to_rgb_float, color_to_rgb_uint8 -# We are the main source, so we use the main_logger -from . import main_logger import torch import torch.nn.functional as F import torchvision.transforms.functional as TF from typing import Optional + +# We are the main source, so we use the main_logger +from . import main_logger +from .helpers import load_image_wrapper try: - from folder_paths import get_input_directory # To get the ComfyUI input directory + from folder_paths import get_input_directory, get_output_directory from comfy import model_management from comfy.utils import common_upscale except ModuleNotFoundError: # No ComfyUI, this is a test environment def get_input_directory(): return "" + get_output_directory = get_input_directory try: from nodes import ImageScale @@ -44,18 +49,13 @@ try: from server import PromptServer except ModuleNotFoundError: PromptServer = None -try: - # We need to import the built-in LoadImage class for ImageDownload - from nodes import LoadImage - has_load_image = True -except Exception: - has_load_image = False logger = main_logger BASE_CATEGORY = "image" IO_CATEGORY = "io" MANIPULATION_CATEGORY = "manipulation" NORMALIZATION = "normalization" +VALIDATION = "validation" FOREGROUND = "foreground" BLUR_SIZE_OPT = ("INT", {"default": 90, "min": 1, "max": 255, "step": 1, }) BLUR_SIZE_TWO_OPT = ("INT", {"default": 6, "min": 1, "max": 255, "step": 1, }) @@ -136,126 +136,281 @@ def parse_size(size_str, reference_dim): return 0 -if has_load_image: - class ImageDownload: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "base_url": ("STRING", { - "default": - "https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/", - "tooltip": "The base URL where the image file is located." - }), - "filename": ("STRING", { - "default": "audioseparation_logo.jpg", - "tooltip": "The name of the image file to download (e.g., photo.jpg, art.png)." - }), - }, - "optional": { - "image_bypass": ("IMAGE", { - "tooltip": "If this image is present will be used instead of the downloaded one" - }), - "mask_bypass": ("MASK", {"tooltip": "If this mask is present will be used instead of the downloaded one"}), - "local_name": ("STRING", { - "default": "", - "tooltip": "The name used locally. Leave empty to use `filename`" - }), - "embed_transparency": ("BOOLEAN", { - "default": False, - "tooltip": "Create RGBA images when they have transparency." - }), - } +# Define sort methods for the node input +sort_methods = [ + "None", + "Alphabetical (ASC)", + "Alphabetical (DESC)", + "Numerical (ASC)", + "Numerical (DESC)", + "Datetime (ASC)", + "Datetime (DESC)" +] + + +# Helper function to extract the first number from a string for sorting +def extract_first_number(s): + match = re.search(r'\d+', s) + return int(match.group()) if match else float('inf') + + +# Sorting function to be used on the lists +def sort_by(items, base_path='.', method=None): + def fullpath(x): return os.path.join(base_path, x) + + def get_timestamp(path): + try: + return os.path.getmtime(path) + except FileNotFoundError: + return float('-inf') + + if method == "Alphabetical (ASC)": + return sorted(items) + elif method == "Alphabetical (DESC)": + return sorted(items, reverse=True) + elif method == "Numerical (ASC)": + return sorted(items, key=lambda x: extract_first_number(os.path.splitext(x)[0])) + elif method == "Numerical (DESC)": + return sorted(items, key=lambda x: extract_first_number(os.path.splitext(x)[0]), reverse=True) + elif method == "Datetime (ASC)": + return sorted(items, key=lambda x: get_timestamp(fullpath(x))) + elif method == "Datetime (DESC)": + return sorted(items, key=lambda x: get_timestamp(fullpath(x)), reverse=True) + else: + return items + + +class ImageDownload: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "base_url": ("STRING", { + "default": + "https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/", + "tooltip": "The base URL where the image file is located." + }), + "filename": ("STRING", { + "default": "audioseparation_logo.jpg", + "tooltip": "The name of the image file to download (e.g., photo.jpg, art.png)." + }), + }, + "optional": { + "image_bypass": ("IMAGE", { + "tooltip": "If this image is present will be used instead of the downloaded one" + }), + "mask_bypass": ("MASK", {"tooltip": "If this mask is present will be used instead of the downloaded one"}), + "local_name": ("STRING", { + "default": "", + "tooltip": "The name used locally. Leave empty to use `filename`" + }), + "embed_transparency": ("BOOLEAN", { + "default": False, + "tooltip": "Create RGBA images when they have transparency." + }), } + } - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("image", "alpha_mask") - FUNCTION = "load_or_download_image" - CATEGORY = BASE_CATEGORY + "/" + IO_CATEGORY - DESCRIPTION = ("Downloads an image to ComfyUI's 'input' directory if it doesn't exist, then loads it using the " - "built-in LoadImage logic.") - UNIQUE_NAME = "SET_ImageDownload" - DISPLAY_NAME = "Image Download and Load" - # This node stores a result to disk. So this IS an output node. - # It can be used without connecting any other node. - # Declaring it as output helps with the preview mechanism. - OUTPUT_NODE = True + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "alpha_mask") + FUNCTION = "load_or_download_image" + CATEGORY = BASE_CATEGORY + "/" + IO_CATEGORY + DESCRIPTION = ("Downloads an image to ComfyUI's 'input' directory if it doesn't exist, then loads it using the " + "built-in LoadImage logic.") + UNIQUE_NAME = "SET_ImageDownload" + DISPLAY_NAME = "Image Download and Load" + # This node stores a result to disk. So this IS an output node. + # It can be used without connecting any other node. + # Declaring it as output helps with the preview mechanism. + OUTPUT_NODE = True - def load_or_download_image(self, base_url: str, filename: str, image_bypass: Optional[torch.Tensor] = None, - mask_bypass: Optional[torch.Tensor] = None, local_name: str = None, - embed_transparency: bool = False): - # If we have something at the bypass inputs use it - if image_bypass is not None or mask_bypass is not None: - if image_bypass is None: - # Just a mask - assert mask_bypass is not None, "This should not be possible if image_bypass is None" # For mypy - image_bypass = torch.zeros(mask_bypass.shape + (3,), dtype=torch.float32, device="cpu") - logger.warning("ImageDownload: Returning an empty image") - elif mask_bypass is None: - # This is ComfyUI behavior when we don't have transparency - mask_bypass = torch.zeros((64, 64), dtype=torch.float32, device="cpu").unsqueeze(0) - logger.warning("ImageDownload: Returning an empty mask") - return (image_bypass, mask_bypass) + def load_or_download_image(self, base_url: str, filename: str, image_bypass: Optional[torch.Tensor] = None, + mask_bypass: Optional[torch.Tensor] = None, local_name: str = None, + embed_transparency: bool = False): + # If we have something at the bypass inputs use it + if image_bypass is not None or mask_bypass is not None: + if image_bypass is None: + # Just a mask + assert mask_bypass is not None, "This should not be possible if image_bypass is None" # For mypy + image_bypass = torch.zeros(mask_bypass.shape + (3,), dtype=torch.float32, device="cpu") + logger.warning("ImageDownload: Returning an empty image") + elif mask_bypass is None: + # This is ComfyUI behavior when we don't have transparency + mask_bypass = torch.zeros((64, 64), dtype=torch.float32, device="cpu").unsqueeze(0) + logger.warning("ImageDownload: Returning an empty mask") + return (image_bypass, mask_bypass) - save_dir = get_input_directory() - dest_fname = local_name or filename - local_filepath = os.path.join(save_dir, dest_fname) + save_dir = get_input_directory() + dest_fname = local_name or filename + local_filepath = os.path.join(save_dir, dest_fname) - if not os.path.exists(local_filepath): - logger.info(f"File '{filename}' not found locally. Attempting to download.") + if not os.path.exists(local_filepath): + logger.info(f"File '{filename}' not found locally. Attempting to download.") - if not base_url.endswith('/'): - base_url += '/' - download_url = base_url + filename + if not base_url.endswith('/'): + base_url += '/' + download_url = base_url + filename - try: - download_file(logger, url=download_url, save_dir=save_dir, file_name=dest_fname, kind="image") - except Exception as e: - logger.error(f"Download failed for {download_url}: {e}", exc_info=True) - raise - else: - logger.info(f"Found existing file, skipping download: '{local_filepath}'") - - # --- REUSE ComfyUI's LoadImage LOGIC --- try: - # Instantiate the built-in LoadImage node - loader_instance = LoadImage() - - # The LoadImage node's `load_image` method expects the filename as passed - # by the ComfyUI widget, which is just the filename. It internally - # resolves the path using folder_paths. - - logger.debug(f"Calling built-in LoadImage.load_image() with filename: '{dest_fname}'") - - # Call the method and return its result directly - result = loader_instance.load_image(dest_fname) - # Create an RGBA image if needed - if embed_transparency: - image, mask = result - # Expand the mask to (b, h, w, 1) - mask = mask[..., None] - # Concatenate image and mask into (b, h, w, 4) - image_with_alpha = torch.cat([image, 1.0 - mask], dim=-1) - result = (image_with_alpha, mask) - # This information is for the preview, as we are an output node and we return images - # they will be displayed in our node. Quite simple. - downloaded_file = { - "images": [{ - "filename": dest_fname, - "subfolder": "", - "type": "input" # We stored the file in the "input" folder - }] - } - return {"ui": downloaded_file, "result": result} - + download_file(logger, url=download_url, save_dir=save_dir, file_name=dest_fname, kind="image") except Exception as e: - logger.error(f"Failed to load image '{filename}' using built-in LoadImage node: {e}", exc_info=True) - # Re-raise to make the error visible in ComfyUI - raise IOError(f"Could not load the image file '{filename}' using the standard loader. " - "It may be corrupt or in an unsupported format.") from e -else: - logger.error("Failed to import ComfyUI `LoadImage`, please fill an issue here: " - "https://github.com/set-soft/ComfyUI-ImageMisc/issues") + logger.error(f"Download failed for {download_url}: {e}", exc_info=True) + raise + else: + logger.info(f"Found existing file, skipping download: '{local_filepath}'") + + return load_image_wrapper(dest_fname, embed_transparency, filename) + + +class ImageLoad: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "file_name": ("STRING", { + "tooltip": "The file name of the image to load" + }), + }, + "optional": { + "embed_transparency": ("BOOLEAN", { + "default": False, + "tooltip": "Create RGBA images when they have transparency." + }), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "alpha_mask") + FUNCTION = "execute" + CATEGORY = BASE_CATEGORY + "/" + IO_CATEGORY + DESCRIPTION = ("Loads an image from any path") + UNIQUE_NAME = "SET_ImageLoad" + DISPLAY_NAME = "Load Image from Path" + # This node stores a result to disk. So this IS an output node. + # It can be used without connecting any other node. + # Declaring it as output helps with the preview mechanism. + OUTPUT_NODE = True + + def execute(self, file_name: str, embed_transparency: bool = False): + if not os.path.exists(file_name): + raise ValueError(f"File '{file_name}' not found") + + return load_image_wrapper(file_name, embed_transparency) + + +class ImageDataset: + """ + A ComfyUI node to prepare lists of images for validation tasks, + such as Salient Object Detection. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "source": ("STRING", { + "default": "./dataset/im", + "tooltip": "Path to the images.\nRelative to ComfyUI input" + }), + "pattern": ("STRING", { + "default": ".*", + "tooltip": "Python regex to match source images." + }), + "destination": ("STRING", { + "default": "./result", + "tooltip": "Path for the result images.\nRelative to ComfyUI output" + }), + "dest_ext": ("STRING", { + "default": "png", + "tooltip": "Extension for the destination images.\nEmpty means same as source" + }), + }, + "optional": { + "reference": ("STRING", { + "default": "./dataset/gt", + "tooltip": "Path for the reference images.\nRelative to ComfyUI input" + }), + "sort_method": (sort_methods,), + } + } + + RETURN_TYPES = ("STRING", "STRING", "STRING",) + RETURN_NAMES = ("images", "results", "references",) + # Tell ComfyUI that the outputs of this node are lists. + OUTPUT_IS_LIST = (True, True, True) + FUNCTION = "generate_lists" + CATEGORY = BASE_CATEGORY + "/" + VALIDATION + UNIQUE_NAME = "SET_ImageDataset" + DISPLAY_NAME = "List Images from Dataset" + + def generate_lists(self, source, pattern, destination, dest_ext, reference=None, sort_method="None"): + # Define valid image extensions + valid_extensions = ['.jpg', '.jpeg', '.png', '.webp'] + source_dir = Path(get_input_directory(), source) + dest_dir = Path(get_output_directory(), destination) + ref_dir = Path(get_input_directory(), reference) if reference else None + + # Ensure directories exist + source_dir.mkdir(parents=True, exist_ok=True) + dest_dir.mkdir(parents=True, exist_ok=True) + if ref_dir: + ref_dir.mkdir(parents=True, exist_ok=True) + + images = [] + results = [] + references = [] + + # Compile the regex pattern + try: + compiled_pattern = re.compile(pattern) + except re.error as e: + raise ValueError(f"Invalid regex pattern: {e}") + + # Get all files in the source directory + source_files = [f for f in os.listdir(source_dir) if (source_dir / f).is_file()] + + # Sort source files before processing + sorted_source_files = sort_by(source_files, base_path=str(source_dir), method=sort_method) + + # Create a lowercase mapping of reference files for case-insensitive matching + ref_map = {} + if ref_dir: + for f in os.listdir(ref_dir): + if (ref_dir / f).is_file(): + ref_map[Path(f).stem.lower()] = f + + for filename in sorted_source_files: + p_filename = Path(filename) + stem = p_filename.stem + ext = p_filename.suffix.lower() + + # Filter by extension and pattern + if ext in valid_extensions and compiled_pattern.search(filename): + # Determine the destination filename and path + dest_extension = f".{dest_ext}" if dest_ext else ext + dest_filename = f"{stem}{dest_extension}" + dest_path = dest_dir / dest_filename + + # Skip if the result file already exists + if dest_path.exists(): + continue + + # Find the reference file (case-insensitive and extension-agnostic) + ref_filename = "" + if ref_dir: + ref_filename_found = ref_map.get(stem.lower()) + if ref_filename_found: + ref_filename = str(ref_dir / ref_filename_found) + + # Add the absolute paths to the lists + images.append(str(source_dir / filename)) + results.append(str(dest_path)) + references.append(ref_filename if ref_dir else "") + + logger.info(f"Found {len(images)} images") + logger.debug(images) + + return (images, results, references) class CompositeFace: