Rename DATASET_SaveImage.py to DataSet_SaveImage.py

2024-07-05 14:22:20 +05:30
parent 5fafc006cc
commit baf9eba4da
1 changed files with 0 additions and 0 deletions
@@ -0,0 +1,75 @@
+import os
+import json
+import numpy as np
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+from comfy.cli_args import args
+
+class DataSet_SaveImage:
+
+    def __init__(self):
+        self.compression = 4
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "Images": ("IMAGE",),
+                "ImageFilePrefix": ("STRING", {"default": "Image"}),
+                "destination": ("STRING", {}),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+    INPUT_IS_LIST = True
+    RETURN_TYPES = ()
+    FUNCTION = "BatchSave"
+    OUTPUT_NODE = True
+    CATEGORY = "🔶DATASET🔶"
+
+    def BatchSave(self, Images, ImageFilePrefix, destination, prompt=None, extra_pnginfo=None):
+        try:
+            ImageFilePrefix = ImageFilePrefix[0]
+            Directory = destination[0]
+
+            if not os.path.exists(Directory):
+                os.makedirs(Directory)
+
+            for i, image in enumerate(Images):
+                image = image.cpu().numpy()
+
+                # Handle image shape to be (H, W, C)
+                if image.ndim == 4:  # If the image has a batch dimension
+                    image = image[0]
+
+                if image.ndim == 3 and image.shape[0] == 1:  # If the image has an extra dimension
+                    image = image[0]
+
+                image = (image * 255).astype(np.uint8)
+                img = Image.fromarray(image)
+                metadata = None
+                if not args.disable_metadata:
+                    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))
+
+                # Generate filename with padding of 4 digits
+                filename = f"{ImageFilePrefix}_{str(i).zfill(5)}.png"
+                file_path = os.path.join(Directory, filename)
+                img.save(file_path, pnginfo=metadata, compress_level=self.compression)
+
+        except Exception as e:
+            print(f"Error saving image: {e}")
+
+        return ()
+
+N_CLASS_MAPPINGS = {
+    "DataSet_SaveImage": DataSet_SaveImage,
+}
+
+N_DISPLAY_NAME_MAPPINGS = {
+    "DataSet_SaveImage": "DataSet_SaveImage",
+}