diff --git a/classes/DataSet_SaveImagePro.py b/classes/DataSet_SaveImagePro.py new file mode 100644 index 0000000..f655f47 --- /dev/null +++ b/classes/DataSet_SaveImagePro.py @@ -0,0 +1,83 @@ +import os +import json +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from comfy.cli_args import args + +def save_images_with_metadata(image_tensors, names, destination, img_format='png', quality=95, disable_metadata=True, extra_pnginfo=None, prompt=None, compression=4): + + if not os.path.exists(destination): + os.makedirs(destination) + + for i, tensor in enumerate(image_tensors): + image = tensor.cpu().numpy() + if image.ndim == 4: + image = image[0] + if image.ndim == 3 and image.shape[0] == 1: + image = image[0] + image = (image * 255).astype(np.uint8) + img = Image.fromarray(image) + metadata = None + if not disable_metadata and img_format.lower() == 'png': + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None and isinstance(extra_pnginfo, dict): + for key, value in extra_pnginfo.items(): + metadata.add_text(key, json.dumps(value)) + + img_name = f"{names[i]}.{img_format}" + img_path = os.path.join(destination, img_name) + + if img_format.lower() in ['jpeg', 'jpg']: + img.save(img_path, format='jpeg', quality=quality) + elif img_format.lower() == 'png': + img.save(img_path, pnginfo=metadata, compress_level=compression) + else: + img.save(img_path, format=img_format.upper()) + + +class DataSet_SaveImagePro: + + def __init__(self): + self.compression = 4 + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "names": ("STRING",), + "destination": ("STRING", {}), + "image_format": (['png', 'jpg'],), + "image_quality": ("INT", {"default": 100, "min": 1, "max": 100}), + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + INPUT_IS_LIST = True + RETURN_TYPES = () + FUNCTION = "BatchSave" + OUTPUT_NODE = True + CATEGORY = "🔶DATASET🔶" + + def BatchSave(self, images, names, destination, image_format, image_quality, prompt=None, extra_pnginfo=None): + + try: + + save_images_with_metadata(images,names,destination[0],image_format[0],image_quality[0],False,extra_pnginfo,prompt,self.compression) + + except Exception as e: + print(f"Error saving image: {e}") + + return () + + +N_CLASS_MAPPINGS = { + "DataSet_SaveImagePro": DataSet_SaveImagePro, +} + +N_DISPLAY_NAME_MAPPINGS = { + "DataSet_SaveImagePro": "DataSet_SaveImagePro", +}