59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
import os
|
|
import random
|
|
import sys
|
|
|
|
# Add custom node directory to path so we can import the class
|
|
sys.path.append(os.path.dirname(__file__))
|
|
|
|
# Mock folder_paths for the node to work
|
|
import folder_paths
|
|
folder_paths.get_annotated_filepath = lambda x: x # Just return the path as is
|
|
|
|
from extract_metadata_node import ImageMetadataExtractor
|
|
|
|
def main():
|
|
output_dir = r"D:\ComfyUI7\ComfyUI\output"
|
|
output_file = "batch_test_results.txt"
|
|
|
|
# Get all image files
|
|
all_files = [os.path.join(output_dir, f) for f in os.listdir(output_dir)
|
|
if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
|
|
|
|
if not all_files:
|
|
print(f"No images found in {output_dir}")
|
|
return
|
|
|
|
# Select 20 random images
|
|
num_samples = min(20, len(all_files))
|
|
selected_files = random.sample(all_files, num_samples)
|
|
|
|
extractor = ImageMetadataExtractor()
|
|
|
|
print(f"Testing on {num_samples} images...")
|
|
|
|
with open(output_file, "w", encoding="utf-8") as f:
|
|
for i, file_path in enumerate(selected_files):
|
|
try:
|
|
# We need to bypass the folder_paths.get_annotated_filepath call inside the node
|
|
# by mocking it, or just passing the absolute path if our mock above works.
|
|
# The node calls folder_paths.get_annotated_filepath(image)
|
|
# Our mock returns x, so we pass the full path.
|
|
|
|
prompt, width, height = extractor.extract_metadata(file_path)
|
|
|
|
separator = "=" * 80
|
|
entry = f"{separator}\nImage: {os.path.basename(file_path)}\nDimensions: {width}x{height}\nPrompt:\n{prompt}\n"
|
|
|
|
f.write(entry + "\n")
|
|
print(f"Processed {i+1}/{num_samples}: {os.path.basename(file_path)}")
|
|
|
|
except Exception as e:
|
|
error_msg = f"Error processing {os.path.basename(file_path)}: {e}\n"
|
|
f.write(error_msg)
|
|
print(error_msg)
|
|
|
|
print(f"Done. Results saved to {output_file}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|