Update fast_saver.py
This commit is contained in:
@@ -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": []}}
|
||||
Reference in New Issue
Block a user