新增 ImageWorkflowExtractor 類別,實現圖像資訊提取與轉換功能

This commit is contained in:
NineKey
2026-01-21 23:35:06 +08:00
parent 2ebceecc53
commit 58dbd81bf1
2 changed files with 102 additions and 2 deletions
+5 -2
View File
@@ -3,6 +3,7 @@ import folder_paths
from .loader_node import DynamicTagLoaderJS
from .saver_node import DynamicTagSaver
from .iterator_node import DynamicTagIterator
from .image_info_node import ImageWorkflowExtractor
# ==============================================================================
# 模組導入與環境檢查
@@ -74,13 +75,15 @@ if PromptServer:
NODE_CLASS_MAPPINGS = {
"DynamicTagLoaderJS": DynamicTagLoaderJS,
"DynamicTagSaver": DynamicTagSaver,
"DynamicTagIterator": DynamicTagIterator
"DynamicTagIterator": DynamicTagIterator,
"WorkflowMetadataReader": ImageWorkflowExtractor,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DynamicTagLoaderJS": "⚡Dynamic Tag Loader",
"DynamicTagSaver": "💾 Dynamic Tag Saver",
"DynamicTagIterator": "🔄 Dynamic Tag Iterator"
"DynamicTagIterator": "🔄 Dynamic Tag Iterator",
"WorkflowMetadataReader": "🔍 Workflow Metadata Reader",
}
WEB_DIRECTORY = "./web"
+97
View File
@@ -0,0 +1,97 @@
import os
import json
import torch
import numpy as np
from PIL import Image, ImageOps
class ImageWorkflowExtractor:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_or_dir": ("STRING", {"default": "C:/ComfyUI/output"}),
"search_by": (["ID", "Type"], {"default": "ID"}),
"search_query": ("STRING", {"default": "71"}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
},
}
# 修改輸出類型:將原本的 raw_info (STRING) 改為 image (IMAGE)
RETURN_TYPES = ("JSON", "STRING", "STRING", "IMAGE")
RETURN_NAMES = ("workflow_json", "clean_text", "selected_path", "image")
FUNCTION = "extract_info"
CATEGORY = "DynamicTags"
def extract_info(self, image_or_dir, search_by, search_query, seed):
target_path = image_or_dir.strip()
selected_file = ""
# ==========================================
# 1. 檔案選取邏輯
# ==========================================
if os.path.isfile(target_path):
selected_file = target_path
elif os.path.isdir(target_path):
valid_extensions = ('.png', '.webp', '.jpg', '.jpeg')
files = sorted([
os.path.join(target_path, f)
for f in os.listdir(target_path)
if f.lower().endswith(valid_extensions)
])
if not files:
# 若無檔案,回傳一個空的黑色張量避免系統崩潰
return ({}, "No images found", target_path, torch.zeros((1, 64, 64, 3)))
index = seed % len(files)
selected_file = files[index]
else:
return ({}, "Invalid path", target_path, torch.zeros((1, 64, 64, 3)))
# ==========================================
# 2. 資訊提取與圖像轉換輸出
# ==========================================
try:
with Image.open(selected_file) as img:
# --- A. 提取 Workflow Metadata ---
info = img.info
workflow = info.get("workflow", "{}")
wf_data = json.loads(workflow) if isinstance(workflow, str) else workflow
nodes = wf_data.get("nodes", [])
clean_results = []
for node in nodes:
match = False
if search_by == "ID":
if str(node.get("id")) == search_query.strip():
match = True
else:
if node.get("type") == search_query.strip():
match = True
if match:
values = node.get("widgets_values", [])
for val in values:
if isinstance(val, list):
for sub_val in val:
clean_results.append(str(sub_val))
else:
clean_results.append(str(val))
final_text = "\n".join(clean_results) if clean_results else f"No nodes match {search_query}"
# --- B. 將圖片轉為 ComfyUI 格式 (IMAGE Tensor) ---
# 修正圖片轉向 (Exif 資訊)
img = ImageOps.exif_transpose(img)
# 統一轉為 RGB
image_rgb = img.convert("RGB")
# 轉為 numpy 陣列並正規化至 0.0 ~ 1.0
image_np = np.array(image_rgb).astype(np.float32) / 255.0
# 轉為 PyTorch Tensor 並調整維度為 [Batch, Height, Width, Channel]
image_tensor = torch.from_numpy(image_np)[None,]
return (wf_data, final_text, selected_file, image_tensor)
except Exception as e:
# 發生錯誤時回傳錯誤訊息與空圖片
return ({}, f"Error: {str(e)}", selected_file, torch.zeros((1, 64, 64, 3)))