diff --git a/__init__.py b/__init__.py index 0401603..df11bdd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,16 +1,19 @@ from .sharp_node import SharpnessAnalyzer, SharpFrameSelector from .parallel_loader import ParallelSharpnessLoader +from .fast_saver import FastAbsoluteSaver # <--- Added this missing import NODE_CLASS_MAPPINGS = { "SharpnessAnalyzer": SharpnessAnalyzer, "SharpFrameSelector": SharpFrameSelector, - "ParallelSharpnessLoader": ParallelSharpnessLoader + "ParallelSharpnessLoader": ParallelSharpnessLoader, + "FastAbsoluteSaver": FastAbsoluteSaver } NODE_DISPLAY_NAME_MAPPINGS = { "SharpnessAnalyzer": "1. Sharpness Analyzer", "SharpFrameSelector": "2. Sharp Frame Selector", - "ParallelSharpnessLoader": "3. Parallel Video Loader (Sharpness)" + "ParallelSharpnessLoader": "3. Parallel Video Loader (Sharpness)", + "FastAbsoluteSaver": "Fast Absolute Saver (Metadata)" } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/fast_saver.py b/fast_saver.py new file mode 100644 index 0000000..d5b05a9 --- /dev/null +++ b/fast_saver.py @@ -0,0 +1,110 @@ +import os +import torch +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +import concurrent.futures +import re + +class FastAbsoluteSaver: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}), + "filename_prefix": ("STRING", {"default": "frame"}), + # NEW: User can define the metadata key name + "metadata_key": ("STRING", {"default": "sharpness_score", "label": "Metadata Key Name"}), + }, + "optional": { + "scores_info": ("STRING", {"forceInput": True}), + } + } + + RETURN_TYPES = () + FUNCTION = "save_images_fast" + OUTPUT_NODE = True + CATEGORY = "BetaHelper/IO" + + def parse_scores(self, scores_str, batch_size): + """ + Parses the string "F:10 (Score:500)..." into a list of floats. + Robust to spaces: handles "Score:500" and "Score: 500" + """ + if not scores_str: + return [0.0] * batch_size + + # Regex explanation: + # Score:\s* -> Matches "Score:" followed by optional spaces + # (\d+(\.\d+)?) -> Matches integer or float (Capture Group 1) + patterns = re.findall(r"Score:\s*(\d+(\.\d+)?)", scores_str) + + scores = [] + for match in patterns: + try: + scores.append(float(match[0])) + except ValueError: + scores.append(0.0) + + # Fill missing scores with 0.0 if batch size mismatches + if len(scores) < batch_size: + scores.extend([0.0] * (batch_size - len(scores))) + return scores[:batch_size] + + def save_single_image(self, tensor_img, full_path, score, key_name): + """Worker function to save one image with metadata""" + try: + array = 255. * tensor_img.cpu().numpy() + img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8)) + + metadata = PngInfo() + + # Use the user-defined key. + # If you want to force no spaces, uncomment the line below: + # key_name = key_name.replace(" ", "_") + + metadata.add_text(key_name, str(score)) + metadata.add_text("software", "ComfyUI_Parallel_Node") + + # compress_level=1 is fast. + img.save(full_path, pnginfo=metadata, compress_level=1) + return True + except Exception as e: + print(f"xx- Error saving {full_path}: {e}") + return False + + def save_images_fast(self, images, output_path, filename_prefix, metadata_key, scores_info=None): + + # 1. Clean Path + output_path = output_path.strip('"') + if not os.path.exists(output_path): + try: + os.makedirs(output_path, exist_ok=True) + except OSError: + raise ValueError(f"Could not create directory: {output_path}") + + # 2. Parse Scores + batch_size = len(images) + scores_list = self.parse_scores(scores_info, batch_size) + + print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...") + + # 3. Parallel Saving + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor: + futures = [] + + for i, img_tensor in enumerate(images): + import time + timestamp = int(time.time()) + # Added index `i` to filename to ensure uniqueness in same batch + fname = f"{filename_prefix}_{timestamp}_{i:03d}.png" + full_path = os.path.join(output_path, fname) + + # Pass the metadata_key to the worker + futures.append(executor.submit(self.save_single_image, img_tensor, full_path, scores_list[i], metadata_key)) + + concurrent.futures.wait(futures) + + print("xx- FastSaver: Save Complete.") + return {"ui": {"images": []}} \ No newline at end of file