From 1bde14bd97e8cb478decb20edbf12333d625a521 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Tue, 20 Jan 2026 00:08:41 +0100 Subject: [PATCH] Update fast_saver.py --- fast_saver.py | 83 ++++++++++++++++++++++++++------------------------- 1 file changed, 43 insertions(+), 40 deletions(-) diff --git a/fast_saver.py b/fast_saver.py index d5b05a9..875fc3c 100644 --- a/fast_saver.py +++ b/fast_saver.py @@ -5,6 +5,7 @@ from PIL import Image from PIL.PngImagePlugin import PngInfo import concurrent.futures import re +import time class FastAbsoluteSaver: @classmethod @@ -14,8 +15,9 @@ class FastAbsoluteSaver: "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"}), + "metadata_key": ("STRING", {"default": "sharpness_score"}), + # NEW: Boolean Switch + "filename_with_score": ("BOOLEAN", {"default": False, "label": "Append Score to Filename"}), }, "optional": { "scores_info": ("STRING", {"forceInput": True}), @@ -27,56 +29,49 @@ class FastAbsoluteSaver: OUTPUT_NODE = True CATEGORY = "BetaHelper/IO" - def parse_scores(self, scores_str, batch_size): + def parse_info(self, info_str, batch_size): """ - Parses the string "F:10 (Score:500)..." into a list of floats. - Robust to spaces: handles "Score:500" and "Score: 500" + Extracts both Frame Indices AND Scores. """ - if not scores_str: - return [0.0] * batch_size + if not info_str: + return ([0]*batch_size, [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) + matches = re.findall(r"F:(\d+).*?Score:\s*(\d+(\.\d+)?)", info_str) + frames = [] 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] + for m in matches: + try: + frames.append(int(m[0])) + scores.append(float(m[1])) + except ValueError: + pass + + if len(frames) < batch_size: + missing = batch_size - len(frames) + frames.extend([0] * missing) + scores.extend([0.0] * missing) + + return frames[:batch_size], 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): + def save_images_fast(self, images, output_path, filename_prefix, metadata_key, filename_with_score, scores_info=None): - # 1. Clean Path output_path = output_path.strip('"') if not os.path.exists(output_path): try: @@ -84,27 +79,35 @@ class FastAbsoluteSaver: 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) + frame_indices, scores_list = self.parse_info(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" + + real_frame_num = frame_indices[i] + current_score = scores_list[i] + + # BASE NAME: frame_001450 + base_name = f"{filename_prefix}_{real_frame_num:06d}" + + # OPTION: Append Score -> frame_001450_1500 + if filename_with_score: + base_name += f"_{int(current_score)}" + + # FALLBACK for missing data + if real_frame_num == 0 and scores_info is None: + base_name = f"{filename_prefix}_{int(time.time())}_{i:03d}" + + fname = f"{base_name}.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)) + futures.append(executor.submit(self.save_single_image, img_tensor, full_path, current_score, metadata_key)) concurrent.futures.wait(futures) - print("xx- FastSaver: Save Complete.") return {"ui": {"images": []}} \ No newline at end of file