diff --git a/__init__.py b/__init__.py index 9a3aeb3..b75b617 100644 --- a/__init__.py +++ b/__init__.py @@ -69,6 +69,8 @@ from .nodes.FL_BulkPDFLoader import FL_BulkPDFLoader from .nodes.FL_SaveAndDisplayImage import FL_SaveAndDisplayImage from .nodes.FL_OllamaCaptioner import FL_OllamaCaptioner from .nodes.FL_ImageAdjuster import FL_ImageAdjuster +from .nodes.FL_Caption_Saver_V2 import FL_CaptionSaver_V2 +from .nodes.FL_PathTypeChecker import FL_PathTypeChecker @@ -145,6 +147,8 @@ NODE_CLASS_MAPPINGS = { "FL_SaveAndDisplayImage": FL_SaveAndDisplayImage, "FL_OllamaCaptioner": FL_OllamaCaptioner, "FL_ImageAdjuster": FL_ImageAdjuster, + "FL_CaptionSaver_V2": FL_CaptionSaver_V2, + "FL_PathTypeChecker": FL_PathTypeChecker, } @@ -220,6 +224,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FL_SaveAndDisplayImage": "FL Save And Display Image", "FL_OllamaCaptioner": "FL Ollama Captioner by Cosmic", "FL_ImageAdjuster": "FL_ImageAdjuster", + "FL_CaptionSaver_V2": "FL Caption Saver V2", + "FL_PathTypeChecker": "FL Path Type Checker", } diff --git a/nodes/FL_Caption_Saver_V2.py b/nodes/FL_Caption_Saver_V2.py new file mode 100644 index 0000000..7846c66 --- /dev/null +++ b/nodes/FL_Caption_Saver_V2.py @@ -0,0 +1,108 @@ +import os +import re +from PIL import Image +import numpy as np +from comfy.utils import ProgressBar + + +class FL_CaptionSaver_V2: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_type": (["Image Input", "Directory Input"], {"default": "Image Input"}), + "caption_input_type": (["Single Caption", "Multiple Captions"], {"default": "Single Caption"}), + "folder_name": ("STRING", {"default": "output_folder"}), + "overwrite": ("BOOLEAN", {"default": True}), + "downsize_factor": ([1, 2, 3], {"default": 1}) + }, + "optional": { + "images": ("IMAGE", {}), + "input_directory": ("STRING", {"default": ""}), + "single_caption": ("STRING", {"default": "Your caption here"}), + "multiple_captions": ("STRING", {"multiline": True, "default": ""}) + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "save_images_with_captions" + CATEGORY = "🏵️Fill Nodes/Captioning" + OUTPUT_NODE = True + + def sanitize_text(self, text): + return re.sub(r'[^a-zA-Z0-9\s.,!?-]', '', text) + + def save_images_with_captions(self, input_type, caption_input_type, folder_name, overwrite, downsize_factor, + images=None, input_directory=None, single_caption="", multiple_captions=""): + os.makedirs(folder_name, exist_ok=True) + + if input_type == "Image Input" and images is not None: + image_list = images + use_original_names = False + elif input_type == "Directory Input" and input_directory: + image_list = [f for f in os.listdir(input_directory) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))] + use_original_names = True + else: + return ("No valid input provided.",) + + if caption_input_type == "Single Caption": + captions = [self.sanitize_text(single_caption)] * len(image_list) + else: + captions = [self.sanitize_text(cap.strip()) for cap in multiple_captions.split('\n') if cap.strip()] + if len(captions) < len(image_list): + captions.extend([captions[-1]] * (len(image_list) - len(captions))) + elif len(captions) > len(image_list): + captions = captions[:len(image_list)] + + saved_files = [] + pbar = ProgressBar(len(image_list)) + for i, (image_item, caption) in enumerate(zip(image_list, captions)): + if use_original_names: + base_name = os.path.splitext(image_item)[0] + image_path = os.path.join(input_directory, image_item) + image = Image.open(image_path) + else: + base_name = f"image_{i}" + image_np = image_item.cpu().numpy() + image_np = self.process_image_tensor(image_np) + image = Image.fromarray(image_np) + + # Downsize the image + if downsize_factor > 1: + new_size = (image.width // downsize_factor, image.height // downsize_factor) + image = image.resize(new_size, Image.LANCZOS) + + image_file_name = f"{folder_name}/{base_name}.png" + text_file_name = f"{folder_name}/{base_name}.txt" + + if not overwrite: + image_file_name, text_file_name = self.get_unique_filenames(folder_name, base_name) + + image.save(image_file_name) + saved_files.append(image_file_name) + + with open(text_file_name, "w") as text_file: + text_file.write(caption) + + pbar.update_absolute(i) + + return (f"Saved {len(saved_files)} images (downsized by factor {downsize_factor}) and captions in '{folder_name}'",) + + def process_image_tensor(self, image_np): + if image_np.shape[0] == 1: + image_np = np.squeeze(image_np, axis=0) + if len(image_np.shape) == 2: + image_np = np.stack((image_np,) * 3, axis=-1) + elif image_np.shape[2] == 1: + image_np = np.repeat(image_np, 3, axis=2) + return (image_np * 255).clip(0, 255).astype(np.uint8) + + def get_unique_filenames(self, folder_name, base_name): + counter = 1 + image_file_name = f"{folder_name}/{base_name}.png" + text_file_name = f"{folder_name}/{base_name}.txt" + while os.path.exists(image_file_name) or os.path.exists(text_file_name): + image_file_name = f"{folder_name}/{base_name}_{counter}.png" + text_file_name = f"{folder_name}/{base_name}_{counter}.txt" + counter += 1 + return image_file_name, text_file_name \ No newline at end of file diff --git a/nodes/FL_Image_Caption_Saver.py b/nodes/FL_Image_Caption_Saver.py index 16911e1..d1dc033 100644 --- a/nodes/FL_Image_Caption_Saver.py +++ b/nodes/FL_Image_Caption_Saver.py @@ -45,17 +45,17 @@ class FL_ImageCaptionSaver: # Convert tensor to numpy array image_np = image_tensor.cpu().numpy() - + # Ensure the image is in the correct shape (height, width, channels) if image_np.shape[0] == 1: # If the first dimension is 1, squeeze it image_np = np.squeeze(image_np, axis=0) - + # If the image is grayscale (2D), convert to RGB if len(image_np.shape) == 2: image_np = np.stack((image_np,) * 3, axis=-1) elif image_np.shape[2] == 1: # If it's (height, width, 1) image_np = np.repeat(image_np, 3, axis=2) - + # Ensure values are in 0-255 range image_np = (image_np * 255).clip(0, 255).astype(np.uint8) diff --git a/nodes/FL_PathTypeChecker.py b/nodes/FL_PathTypeChecker.py new file mode 100644 index 0000000..ae71f1e --- /dev/null +++ b/nodes/FL_PathTypeChecker.py @@ -0,0 +1,45 @@ +import os +import pathlib + +class FL_PathTypeChecker: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_path": ("STRING", {"default": "", "multiline": False}), + } + } + + RETURN_TYPES = ("PATH",) + FUNCTION = "check_path_type" + CATEGORY = "🏵️Fill Nodes/Utils" + + def check_path_type(self, input_path): + input_path = input_path.strip() # Remove leading/trailing whitespace + + if not input_path: + return ("Empty path provided.",) + + path = pathlib.Path(input_path) + + if path.is_absolute(): + return ("Absolute path",) + elif path.is_relative_to(pathlib.Path.cwd()): + return ("Relative path",) + else: + # Check if it might be a valid relative path + try: + path.relative_to(".") + return ("Relative path",) + except ValueError: + pass + + # If it's not recognized as absolute or relative, it might be invalid or a special case + if os.path.splitdrive(input_path)[0]: + return ("Drive-specific path",) + elif input_path.startswith('//') or input_path.startswith('\\\\'): + return ("UNC path",) + elif '://' in input_path: + return ("URL-like path",) + else: + return ("Unrecognized or invalid path",) \ No newline at end of file