Compare commits
18 Commits
1bde14bd97
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| f7c4acfebb | |||
| 6c7f618bb0 | |||
| f02851d88a | |||
| 59bd805920 | |||
| 3297f3a203 | |||
| 3b87ac820f | |||
| b23773c7c2 | |||
| 8560f24d36 | |||
| c40c1fd82c | |||
| 07acefffc1 | |||
| b30e2d0233 | |||
| dd6a9aefd7 | |||
| 60362e3514 | |||
| f03979a767 | |||
| c32a4bcb32 | |||
| 162699a4a2 | |||
| f63b837a2c | |||
| b54b4329ca |
1
.gitignore
vendored
Normal file
1
.gitignore
vendored
Normal file
@@ -0,0 +1 @@
|
||||
ffmpeg_bin/
|
||||
134
README.md
134
README.md
@@ -1,24 +1,30 @@
|
||||
# 🔪 ComfyUI Sharp Frame Selector
|
||||
# ComfyUI Sharpness Helper Nodes
|
||||
|
||||
A collection of high-performance custom nodes for **ComfyUI** designed to detect blur, calculate sharpness scores, and automatically extract the best frames from videos or image batches.
|
||||
A high-performance custom node suite for **ComfyUI** designed to detect blur, calculate sharpness scores (Laplacian Variance), and efficiently extract or filter the best frames from videos and image batches.
|
||||
|
||||
This pack includes two distinct approaches:
|
||||
1. **Parallel Video Loader:** A multi-threaded, path-based loader for processing massive video files directly from disk (Low RAM usage).
|
||||
2. **Standard Sharpness Duo:** A classic filter setup for processing images/latents *inside* your existing workflow.
|
||||
This pack is for a personnal project:
|
||||
1. **Dataset Creation:** Extracting only the sharpest frames from massive movie files without crashing RAM.
|
||||
2. **Generation Filtering:** Automatically discarding blurry frames from Wan or img2img outputs.
|
||||
|
||||
---
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Key Features
|
||||
|
||||
### 1. New: Parallel Video Loader (Path-Based)
|
||||
### 1. Parallel Video Loader (Path-Based)
|
||||
* **Zero-RAM Scanning:** Scans video files directly from disk without decoding every frame to memory.
|
||||
* **Multi-Threaded:** Uses all CPU cores to calculate sharpness scores at high speed.
|
||||
* **Batching Support:** Includes a "Page" system to process long movies in chunks (e.g., minute-by-minute) without restarting ComfyUI.
|
||||
* **Smart Selection:** Automatically skips "adjacent" frames to ensure you get a diverse selection of sharp images.
|
||||
* **Multi-Threaded:** Uses all CPU cores to calculate sharpness scores at high speed (1000s of frames per minute).
|
||||
* **Smart Batching:** Includes an auto-incrementing "Page" system to process long movies in chunks (e.g., minute-by-minute) without restarting ComfyUI.
|
||||
* **Lazy Loading:** Only decodes and loads the final "Best N" frames into ComfyUI tensors.
|
||||
|
||||
### 2. Standard Sharpness Duo (Tensor-Based)
|
||||
### 2. Fast Absolute Saver (Metadata)
|
||||
* **Multi-Threaded Saving:** Spawns parallel workers to saturate SSD write speeds (bypassing standard PIL bottlenecks).
|
||||
* **No UI Lag:** Saves images in the background without trying to render Base64 previews in the browser, preventing interface freezes.
|
||||
* **Metadata Embedding:** Automatically embeds the sharpness score into the PNG/WebP metadata for dataset curation.
|
||||
* **Smart Naming:** Uses original video frame numbers in filenames (e.g., `frame_001450.png`) instead of arbitrary counters.
|
||||
|
||||
### 3. Standard Sharpness Duo (Tensor-Based)
|
||||
* **Workflow Integration:** Works with any node that outputs an `IMAGE` batch (e.g., AnimateDiff, VideoHelperSuite).
|
||||
* **Precision Filtering:** Sorts and filters generated frames before saving or passing to a second pass (img2img).
|
||||
|
||||
@@ -29,4 +35,108 @@ This pack includes two distinct approaches:
|
||||
1. Clone this repository into your `custom_nodes` folder:
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/
|
||||
git clone https://github.com/ethanfel/ComfyUI-Sharp-Selector.git
|
||||
git clone [https://github.com/YOUR_USERNAME/ComfyUI-Sharpness-Helper.git](https://github.com/YOUR_USERNAME/ComfyUI-Sharpness-Helper.git)
|
||||
```
|
||||
2. Install dependencies (if needed):
|
||||
```bash
|
||||
pip install opencv-python numpy
|
||||
```
|
||||
3. Restart ComfyUI.
|
||||
|
||||
---
|
||||
|
||||
## 🛠️ Node Documentation
|
||||
|
||||
### 1. Parallel Video Loader (Sharpness)
|
||||
**Category:** `BetaHelper/Video`
|
||||
|
||||
This is the recommended node for **Dataset Creation** or finding good frames in **Long Movies**. It inputs a file path, scans it in parallel, and only loads the final "Best N" frames into memory.
|
||||
|
||||
| Input | Description |
|
||||
| :--- | :--- |
|
||||
| **video_path** | Absolute path to your video file (e.g., `D:\Movies\input.mp4`). |
|
||||
| **batch_index** | **Critical.** Connect a **Primitive Node** here set to `increment`. This controls which "chunk" of the video you are viewing. |
|
||||
| **scan_limit** | How many frames to process per batch (e.g., `1440`). |
|
||||
| **frame_scan_step** | Speed up scanning by checking every Nth frame (e.g., `5` checks frames 0, 5, 10...). |
|
||||
| **manual_skip_start** | Global offset (e.g., set to `2000` to always ignore the opening credits). |
|
||||
|
||||
**Outputs:**
|
||||
* `images`: The batch of the sharpest frames found.
|
||||
* `scores_info`: String containing frame indices and scores (Connect to Saver).
|
||||
* `batch_int`: The current batch number.
|
||||
* `batch_status`: Human-readable status (e.g., *"Batch 2: Skipped 2880 frames..."*).
|
||||
|
||||
> **💡 Pro Tip:** To scan a movie continuously, connect a **Primitive Node** to `batch_index`, set it to **increment**, and enable "Auto Queue" in ComfyUI.
|
||||
|
||||
---
|
||||
|
||||
### 2. Fast Absolute Saver (Metadata)
|
||||
**Category:** `BetaHelper/IO`
|
||||
|
||||
A "Pro-Grade" saver designed for speed. It bypasses relative paths and UI previews.
|
||||
|
||||
| Input | Description |
|
||||
| :--- | :--- |
|
||||
| **output_path** | Absolute path to save folder (e.g., `D:\Datasets\Sharp_Output`). |
|
||||
| **filename_prefix** | Base name for files (e.g., `matrix_movie`). |
|
||||
| **max_threads** | **0 = Auto** (Uses all CPU cores). Set manually to limit CPU usage. |
|
||||
| **save_format** | `png` (Fastest) or `webp` (Smaller size). |
|
||||
| **filename_with_score** | If True, appends score to filename: `frame_001450_1500.png`. |
|
||||
| **scores_info** | Connect this to the `scores_info` output of the Parallel Loader to enable smart naming. |
|
||||
|
||||
**Performance Note:**
|
||||
* **PNG:** Uses `compress_level=1` for maximum speed.
|
||||
* **WebP:** Avoid `webp_method=6` unless you need max compression; it is very CPU intensive. `4` is the recommended balance.
|
||||
|
||||
---
|
||||
|
||||
### 3. Sharpness Analyzer & Selector (The Duo)
|
||||
**Category:** `BetaHelper/Image`
|
||||
|
||||
Use these when you already have images inside your workflow (e.g., from a generation or a standard Load Video node).
|
||||
|
||||
#### Node A: Sharpness Analyzer
|
||||
* **Input:** `IMAGE` batch.
|
||||
* **Action:** Calculates the Laplacian Variance for every image in the batch.
|
||||
* **Output:** Passes the images through + a generic score list.
|
||||
|
||||
#### Node B: SharpFrame Selector
|
||||
* **Input:** `IMAGE` batch (from Analyzer).
|
||||
* **Action:** Sorts the batch based on the scores and picks the top N frames.
|
||||
* **Output:** A reduced batch containing only the sharpest images.
|
||||
|
||||
---
|
||||
|
||||
## ⚖️ Which Node Should I Use?
|
||||
|
||||
| Feature | **Parallel Video Loader** | **Standard Duo** |
|
||||
| :--- | :--- | :--- |
|
||||
| **Input Type** | File Path (`String`) | Image Tensor (`IMAGE`) |
|
||||
| **Best For** | **Long Videos / Movies** | **Generations / Short Clips** |
|
||||
| **Memory Usage** | Very Low (Only loads final frames) | High (Loads all frames to RAM first) |
|
||||
| **Speed** | ⚡ **Ultra Fast** (Multi-core) | 🐢 Standard (Single-core) |
|
||||
| **Workflow Stage** | Start of Workflow | Middle/End of Workflow |
|
||||
|
||||
---
|
||||
|
||||
## 📝 Example Workflows
|
||||
|
||||
### Batch Processing a Movie for Training Data
|
||||
1. Add **Parallel Video Loader**.
|
||||
2. Connect a **Primitive Node** to `batch_index` (Control: `increment`).
|
||||
3. Set `scan_limit` to `1000` and `frame_scan_step` to `5`.
|
||||
4. Connect `images` and `scores_info` to **Fast Absolute Saver**.
|
||||
5. Enable **Auto Queue** in ComfyUI extra options.
|
||||
* *Result: ComfyUI will loop through your movie, extracting the 4 sharpest frames from every ~1000 frame chunk automatically.*
|
||||
|
||||
### Filtering AnimateDiff Output
|
||||
1. AnimateDiff Generation -> **Sharpness Analyzer**.
|
||||
2. Analyzer Output -> **SharpFrame Selector** (Select Best 1).
|
||||
3. Selector Output -> **Face Detailer** or **Upscaler**.
|
||||
* *Result: Only the clearest frame from your animation is sent to the upscaler, saving time on blurry frames.*
|
||||
|
||||
---
|
||||
|
||||
## Credits
|
||||
* Built using `opencv-python` for Laplacian Variance calculation.
|
||||
* Parallel processing logic for efficient large-file handling.
|
||||
@@ -1,19 +1,16 @@
|
||||
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,
|
||||
"FastAbsoluteSaver": FastAbsoluteSaver
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SharpnessAnalyzer": "1. Sharpness Analyzer",
|
||||
"SharpFrameSelector": "2. Sharp Frame Selector",
|
||||
"ParallelSharpnessLoader": "3. Parallel Video Loader (Sharpness)",
|
||||
"FastAbsoluteSaver": "Fast Absolute Saver (Metadata)"
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
BIN
assets/nodes.png
BIN
assets/nodes.png
Binary file not shown.
|
Before Width: | Height: | Size: 169 KiB After Width: | Height: | Size: 347 KiB |
@@ -1,8 +1,8 @@
|
||||
{
|
||||
"id": "4fbf6f31-0f7b-4465-8ec8-25df4862e076",
|
||||
"revision": 0,
|
||||
"last_node_id": 34,
|
||||
"last_link_id": 42,
|
||||
"last_node_id": 35,
|
||||
"last_link_id": 44,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 31,
|
||||
@@ -33,69 +33,7 @@
|
||||
"version": "7.5.2",
|
||||
"input_ue_unconnectable": {}
|
||||
}
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 29,
|
||||
"type": "ParallelSharpnessLoader",
|
||||
"pos": [
|
||||
4992,
|
||||
-1024
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
262
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
42
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "scores_info",
|
||||
"type": "STRING",
|
||||
"links": []
|
||||
},
|
||||
{
|
||||
"name": "batch_int",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "batch_status",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
39
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"aux_id": "ComfyUI-Sharp-Selector.git",
|
||||
"ver": "dab38a1fbf0077655fe568d500866fce6ecc857d",
|
||||
"Node name for S&R": "ParallelSharpnessLoader",
|
||||
"ue_properties": {
|
||||
"widget_ue_connectable": {},
|
||||
"input_ue_unconnectable": {},
|
||||
"version": "7.5.2"
|
||||
}
|
||||
},
|
||||
"widgets_values": [
|
||||
"C:\\path\\to\\video.mp4",
|
||||
0,
|
||||
1440,
|
||||
1,
|
||||
30,
|
||||
24,
|
||||
2000
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
@@ -109,7 +47,7 @@
|
||||
174
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
@@ -155,64 +93,6 @@
|
||||
0
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 34,
|
||||
"type": "Note",
|
||||
"pos": [
|
||||
4224,
|
||||
-1120
|
||||
],
|
||||
"size": [
|
||||
416,
|
||||
736
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"📌 ComfyUI Sharpness Tools Explained\n\n1. Parallel Video Loader (The \"Loader\")\n\n Best for: Processing long videos or movie files directly from disk.\n\n How it works: It opens the video file itself and uses multi-threading (parallel CPU cores) to scan thousands of frames without loading them into memory. It only \"decodes\" the final few sharpest frames.\n\n Use Case: Extracting dataset images, finding high-quality frames from a raw movie file, or scanning 10,000 frames without crashing your RAM.\n\n Key Feature: Features a \"Batch Counter\" to automatically page through long videos (e.g., scan minute 0-1, then minute 1-2).\n\n2. Standard Sharpness Duo (The \"Filter\")\n\n Best for: Processing images already inside your workflow (e.g., after an img2img pass, or a short generated GIF).\n\n How it works:\n\n Node A (Analyzer): Assigns a score to every image in the batch.\n\n Node B (Selector): Picks the best ones based on those scores.\n\n Use Case: Filtering bad generations, picking the best frame from a small batch of AnimateDiff results, or cleaning up a sequence.\n\n Limitation: It is single-threaded and requires all images to be loaded in VRAM/RAM first (slow for long videos).\n\n🚀 Which one to use?\n\n Starting from a Video File? → Use Parallel Loader.\n\n Starting from a Generation/Latent? → Use Standard Duo."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 33,
|
||||
"type": "SaveImage",
|
||||
"pos": [
|
||||
5824,
|
||||
-1024
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 42
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.9.2",
|
||||
"ue_properties": {
|
||||
"widget_ue_connectable": {},
|
||||
"input_ue_unconnectable": {}
|
||||
},
|
||||
"Node name for S&R": "SaveImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"sharp/img_"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 32,
|
||||
"type": "easy showAnything",
|
||||
@@ -225,7 +105,7 @@
|
||||
96
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
@@ -268,7 +148,7 @@
|
||||
26
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
@@ -298,6 +178,147 @@
|
||||
}
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 35,
|
||||
"type": "FastAbsoluteSaver",
|
||||
"pos": [
|
||||
5856,
|
||||
-1024
|
||||
],
|
||||
"size": [
|
||||
306.3776153564453,
|
||||
270
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 43
|
||||
},
|
||||
{
|
||||
"name": "scores_info",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"link": 44
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"aux_id": "ComfyUI-Sharp-Selector.git",
|
||||
"ver": "162699a4a23219ac5ac75f398a17e67c3767da46",
|
||||
"ue_properties": {
|
||||
"widget_ue_connectable": {},
|
||||
"input_ue_unconnectable": {}
|
||||
},
|
||||
"Node name for S&R": "FastAbsoluteSaver"
|
||||
},
|
||||
"widgets_values": [
|
||||
"D:\\Datasets\\Sharp_Output",
|
||||
"frame",
|
||||
"png",
|
||||
0,
|
||||
false,
|
||||
"sharpness_score",
|
||||
true,
|
||||
100,
|
||||
4
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 29,
|
||||
"type": "ParallelSharpnessLoader",
|
||||
"pos": [
|
||||
4992,
|
||||
-1024
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
262
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
43
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "scores_info",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
44
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "batch_int",
|
||||
"type": "INT",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"name": "batch_status",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
39
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"aux_id": "ComfyUI-Sharp-Selector.git",
|
||||
"ver": "dab38a1fbf0077655fe568d500866fce6ecc857d",
|
||||
"Node name for S&R": "ParallelSharpnessLoader",
|
||||
"ue_properties": {
|
||||
"widget_ue_connectable": {},
|
||||
"input_ue_unconnectable": {},
|
||||
"version": "7.5.2"
|
||||
}
|
||||
},
|
||||
"widgets_values": [
|
||||
"C:\\path\\to\\video.mp4",
|
||||
0,
|
||||
1440,
|
||||
1,
|
||||
30,
|
||||
24,
|
||||
2000
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 34,
|
||||
"type": "Note",
|
||||
"pos": [
|
||||
4224,
|
||||
-1120
|
||||
],
|
||||
"size": [
|
||||
416,
|
||||
736
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"ue_properties": {
|
||||
"widget_ue_connectable": {},
|
||||
"version": "7.5.2",
|
||||
"input_ue_unconnectable": {}
|
||||
}
|
||||
},
|
||||
"widgets_values": [
|
||||
"📝 Smart Dataset Extraction Workflow\n\n1. Parallel Video Loader (The Source)\n\n What it does: Scans your video file directly from the hard drive using multi-threading. It does not load the whole video into RAM.\n\n Batching: Uses the batch_index (Primitive Node) to \"page\" through the movie.\n\n Example: If scan_limit is 1440, Batch 0 scans frames 0-1440, Batch 1 scans 1440-2880, etc.\n\n Selection: It calculates sharpness (Laplacian Variance) and only decodes the \"Best N\" frames to send downstream.\n\n2. Fast Absolute Saver (The Destination)\n\n What it does: Saves images instantly to your SSD using parallel workers, bypassing the slow ComfyUI preview window.\n\n Smart Naming: Connect scores_info from the Loader to the Saver! This allows files to be named using the original video frame number (e.g., movie_frame_00450.png) rather than a random batch counter.\n\n Metadata: Embeds the sharpness score into the PNG/WebP metadata for future filtering.\n\n⚠️ Usage Tips:\n\n Automation: Set batch_index to \"Increment\" (on the Primitive Node) and enable \"Auto Queue\" in ComfyUI options to process the entire movie automatically.\n\n Monitoring: Watch the Console Window (black command prompt) for progress logs. The saver does not preview images in the UI to prevent lag.\n\n Safety: The saver uses absolute paths and overwrites files with the same name. Use a unique filename_prefix for each new video source."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
@@ -318,12 +339,20 @@
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
42,
|
||||
43,
|
||||
29,
|
||||
0,
|
||||
33,
|
||||
35,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
44,
|
||||
29,
|
||||
1,
|
||||
35,
|
||||
1,
|
||||
"STRING"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
@@ -333,10 +362,10 @@
|
||||
"ue_links": [],
|
||||
"links_added_by_ue": [],
|
||||
"ds": {
|
||||
"scale": 0.8264462809917354,
|
||||
"scale": 1.1,
|
||||
"offset": [
|
||||
-2054.6525613149634,
|
||||
1737.5186871750661
|
||||
-3048.698587382934,
|
||||
1363.985488079904
|
||||
]
|
||||
},
|
||||
"frontendVersion": "1.36.14",
|
||||
|
||||
113
fast_saver.py
113
fast_saver.py
@@ -1,113 +0,0 @@
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
import concurrent.futures
|
||||
import re
|
||||
import time
|
||||
|
||||
class FastAbsoluteSaver:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", ),
|
||||
"output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}),
|
||||
"filename_prefix": ("STRING", {"default": "frame"}),
|
||||
"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}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_images_fast"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "BetaHelper/IO"
|
||||
|
||||
def parse_info(self, info_str, batch_size):
|
||||
"""
|
||||
Extracts both Frame Indices AND Scores.
|
||||
"""
|
||||
if not info_str:
|
||||
return ([0]*batch_size, [0.0]*batch_size)
|
||||
|
||||
matches = re.findall(r"F:(\d+).*?Score:\s*(\d+(\.\d+)?)", info_str)
|
||||
|
||||
frames = []
|
||||
scores = []
|
||||
|
||||
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):
|
||||
try:
|
||||
array = 255. * tensor_img.cpu().numpy()
|
||||
img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
|
||||
|
||||
metadata = PngInfo()
|
||||
metadata.add_text(key_name, str(score))
|
||||
metadata.add_text("software", "ComfyUI_Parallel_Node")
|
||||
|
||||
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, filename_with_score, scores_info=None):
|
||||
|
||||
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}")
|
||||
|
||||
batch_size = len(images)
|
||||
frame_indices, scores_list = self.parse_info(scores_info, batch_size)
|
||||
|
||||
print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...")
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
|
||||
futures = []
|
||||
|
||||
for i, img_tensor in enumerate(images):
|
||||
|
||||
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)
|
||||
|
||||
futures.append(executor.submit(self.save_single_image, img_tensor, full_path, current_score, metadata_key))
|
||||
|
||||
concurrent.futures.wait(futures)
|
||||
|
||||
return {"ui": {"images": []}}
|
||||
Reference in New Issue
Block a user