Add folder image loader node
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
@@ -22,6 +23,17 @@ _H3_PLAN_TYPE = "H3_CHAIN_PLAN"
|
||||
_H3_PLAN_IMAGE_BINDINGS_KEY = "_dumas_scene_image_bindings"
|
||||
_H3_PLAN_IMAGE_BINDINGS = OrderedDict()
|
||||
_H3_PLAN_IMAGE_BINDINGS_CAP = 128
|
||||
_FOLDER_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tiff", ".tif")
|
||||
_LOAD_IMAGES_FOLDER_DEFAULT_STATE = {
|
||||
"version": 1,
|
||||
"folder": "",
|
||||
"recursive": False,
|
||||
"sort": "name",
|
||||
"sort_dir": "asc",
|
||||
"selected": [],
|
||||
"selection_mode": "selected",
|
||||
"first_n": 5,
|
||||
}
|
||||
|
||||
|
||||
def _clean_input_token_value(value):
|
||||
@@ -255,6 +267,155 @@ def _tensor_image_to_pil_image(tensor):
|
||||
return Image.fromarray(np.clip(image_array, 0, 255).astype(np.uint8))
|
||||
|
||||
|
||||
def _folder_loader_default_state():
|
||||
return dict(_LOAD_IMAGES_FOLDER_DEFAULT_STATE)
|
||||
|
||||
|
||||
def _parse_load_images_folder_state(state_json):
|
||||
if not state_json:
|
||||
return _folder_loader_default_state()
|
||||
try:
|
||||
parsed = json.loads(state_json)
|
||||
except Exception:
|
||||
return _folder_loader_default_state()
|
||||
|
||||
state = _folder_loader_default_state()
|
||||
if isinstance(parsed, dict):
|
||||
state.update({key: value for key, value in parsed.items() if key in state})
|
||||
return state
|
||||
|
||||
|
||||
def _folder_is_image(name):
|
||||
return str(name or "").lower().endswith(_FOLDER_IMAGE_EXTS)
|
||||
|
||||
|
||||
def _list_folder_image_files(real_folder, recursive):
|
||||
files = []
|
||||
if recursive:
|
||||
for root, _dirs, names in os.walk(real_folder):
|
||||
for name in names:
|
||||
if not _folder_is_image(name):
|
||||
continue
|
||||
full_path = os.path.join(root, name)
|
||||
try:
|
||||
stat_result = os.stat(full_path)
|
||||
except OSError:
|
||||
continue
|
||||
rel_path = os.path.relpath(full_path, real_folder).replace("\\", "/")
|
||||
files.append(
|
||||
{
|
||||
"file": rel_path,
|
||||
"name": name,
|
||||
"size": stat_result.st_size,
|
||||
"mtime": stat_result.st_mtime,
|
||||
}
|
||||
)
|
||||
else:
|
||||
for name in os.listdir(real_folder):
|
||||
full_path = os.path.join(real_folder, name)
|
||||
if not os.path.isfile(full_path) or not _folder_is_image(name):
|
||||
continue
|
||||
try:
|
||||
stat_result = os.stat(full_path)
|
||||
except OSError:
|
||||
continue
|
||||
files.append(
|
||||
{
|
||||
"file": name,
|
||||
"name": name,
|
||||
"size": stat_result.st_size,
|
||||
"mtime": stat_result.st_mtime,
|
||||
}
|
||||
)
|
||||
return files
|
||||
|
||||
|
||||
def _sort_folder_image_files(files, sort_key, sort_dir):
|
||||
ordered = list(files or [])
|
||||
|
||||
def sort_value(entry):
|
||||
if sort_key == "date":
|
||||
return (float(entry.get("mtime") or 0), str(entry.get("file") or "").lower())
|
||||
return str(entry.get("file") or "").lower()
|
||||
|
||||
ordered.sort(key=sort_value, reverse=str(sort_dir or "").lower() == "desc")
|
||||
return ordered
|
||||
|
||||
|
||||
def _resolve_folder_selection(state, files):
|
||||
ordered = _sort_folder_image_files(
|
||||
files,
|
||||
state.get("sort", "name"),
|
||||
state.get("sort_dir", "asc"),
|
||||
)
|
||||
mode = str(state.get("selection_mode") or "selected").lower()
|
||||
if mode == "all":
|
||||
return [entry["file"] for entry in ordered]
|
||||
if mode == "first_n":
|
||||
try:
|
||||
count = max(0, int(state.get("first_n", 0) or 0))
|
||||
except Exception:
|
||||
count = 0
|
||||
return [entry["file"] for entry in ordered[:count]]
|
||||
if mode == "random":
|
||||
return [random.choice(ordered)["file"]] if ordered else []
|
||||
|
||||
present = {entry["file"] for entry in ordered}
|
||||
selected = []
|
||||
for rel_path in state.get("selected", []) or []:
|
||||
if isinstance(rel_path, str) and rel_path in present:
|
||||
selected.append(rel_path)
|
||||
return selected
|
||||
|
||||
|
||||
def _load_folder_image(path):
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
import torch
|
||||
except Exception as exc:
|
||||
raise RuntimeError("torch is required to load folder images") from exc
|
||||
|
||||
from PIL import Image, ImageOps, ImageSequence
|
||||
|
||||
try:
|
||||
import comfy.model_management as comfy_model_management
|
||||
|
||||
tensor_dtype = comfy_model_management.intermediate_dtype()
|
||||
except Exception:
|
||||
tensor_dtype = torch.float32
|
||||
|
||||
try:
|
||||
import node_helpers
|
||||
|
||||
image = node_helpers.pillow(Image.open, path)
|
||||
except Exception:
|
||||
image = Image.open(path)
|
||||
|
||||
frame = ImageOps.exif_transpose(next(ImageSequence.Iterator(image)))
|
||||
if frame.mode == "I":
|
||||
frame = frame.point(lambda px: px * (1 / 255))
|
||||
rgb_image = frame.convert("RGB")
|
||||
width, height = rgb_image.size
|
||||
|
||||
if "A" in frame.getbands():
|
||||
alpha = np.array(frame.getchannel("A")).astype(np.float32) / 255.0
|
||||
mask_image = Image.fromarray(((1.0 - alpha) * 255).astype(np.uint8), mode="L")
|
||||
elif frame.mode == "P" and "transparency" in frame.info:
|
||||
alpha = np.array(frame.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0
|
||||
mask_image = Image.fromarray(((1.0 - alpha) * 255).astype(np.uint8), mode="L")
|
||||
else:
|
||||
mask_image = Image.new("L", rgb_image.size, 0)
|
||||
|
||||
image_tensor = torch.from_numpy(np.array(rgb_image).astype(np.float32) / 255.0)[None,].to(
|
||||
dtype=tensor_dtype
|
||||
)
|
||||
mask_tensor = torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0).to(
|
||||
dtype=tensor_dtype
|
||||
)
|
||||
return image_tensor, mask_tensor, int(width), int(height)
|
||||
|
||||
|
||||
def _normalize_free_text(value):
|
||||
return " ".join(str(value or "").split()).strip()
|
||||
|
||||
@@ -628,6 +789,122 @@ class DumasSaveImageNode:
|
||||
return {"ui": {"images": ui_images}}
|
||||
|
||||
|
||||
class DumasLoadImagesFolderNode:
|
||||
DESCRIPTION = (
|
||||
"Load many images from any folder on disk and feed them through your "
|
||||
"workflow one at a time. Pick specific images, all images, the first N "
|
||||
"images in sort order, or one random image per run."
|
||||
)
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING", "INT", "INT")
|
||||
RETURN_NAMES = ("image", "mask", "width", "height", "filename", "index", "total")
|
||||
OUTPUT_IS_LIST = (True, True, True, True, True, True, True)
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "Dumas/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {},
|
||||
"hidden": {
|
||||
"LoadImagesFolderState": (
|
||||
"STRING",
|
||||
{"default": json.dumps(_LOAD_IMAGES_FOLDER_DEFAULT_STATE)},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def load(self, LoadImagesFolderState=""):
|
||||
state = _parse_load_images_folder_state(LoadImagesFolderState)
|
||||
folder = str(state.get("folder") or "").strip()
|
||||
recursive = bool(state.get("recursive", False))
|
||||
|
||||
if not folder or not os.path.isdir(folder):
|
||||
raise ValueError(
|
||||
"Load Images from Folder: folder not found. Set a folder on the node first."
|
||||
)
|
||||
|
||||
real_folder = os.path.realpath(folder)
|
||||
files = _list_folder_image_files(real_folder, recursive)
|
||||
selected = _resolve_folder_selection(state, files)
|
||||
mode = str(state.get("selection_mode") or "selected").lower()
|
||||
if not selected:
|
||||
if mode == "random":
|
||||
raise ValueError(
|
||||
"Load Images from Folder: no images found for random selection."
|
||||
)
|
||||
raise ValueError(
|
||||
"Load Images from Folder: no images selected. Use Pick images on the node."
|
||||
)
|
||||
|
||||
images = []
|
||||
masks = []
|
||||
widths = []
|
||||
heights = []
|
||||
names = []
|
||||
indices = []
|
||||
|
||||
count = 0
|
||||
for rel_path in selected:
|
||||
if not isinstance(rel_path, str) or not rel_path:
|
||||
continue
|
||||
full_path = os.path.realpath(os.path.join(real_folder, rel_path))
|
||||
if not _is_within_directory(real_folder, full_path) or not os.path.isfile(full_path):
|
||||
continue
|
||||
try:
|
||||
image_tensor, mask_tensor, width, height = _load_folder_image(full_path)
|
||||
except Exception as exc:
|
||||
print(f"[DumasLoadImagesFolder] failed to load {rel_path}: {exc}")
|
||||
continue
|
||||
|
||||
images.append(image_tensor)
|
||||
masks.append(mask_tensor)
|
||||
widths.append(width)
|
||||
heights.append(height)
|
||||
if recursive:
|
||||
names.append(os.path.splitext(rel_path)[0].replace("/", "_").replace("\\", "_"))
|
||||
else:
|
||||
names.append(os.path.splitext(os.path.basename(rel_path))[0])
|
||||
count += 1
|
||||
indices.append(count)
|
||||
|
||||
if not images:
|
||||
raise ValueError(
|
||||
"Load Images from Folder: none of the chosen images could be loaded."
|
||||
)
|
||||
|
||||
totals = [count] * len(images)
|
||||
return (images, masks, widths, heights, names, indices, totals)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, LoadImagesFolderState=""):
|
||||
state = _parse_load_images_folder_state(LoadImagesFolderState)
|
||||
folder = str(state.get("folder") or "").strip()
|
||||
if not folder or not os.path.isdir(folder):
|
||||
return hashlib.sha256((LoadImagesFolderState or "").encode("utf-8")).hexdigest()
|
||||
|
||||
real_folder = os.path.realpath(folder)
|
||||
files = _list_folder_image_files(real_folder, bool(state.get("recursive", False)))
|
||||
mode = str(state.get("selection_mode") or "selected").lower()
|
||||
parts = [json.dumps({k: v for k, v in state.items() if k != "selected"}, sort_keys=True)]
|
||||
|
||||
if mode == "random":
|
||||
for entry in _sort_folder_image_files(files, state.get("sort", "name"), state.get("sort_dir", "asc")):
|
||||
parts.append(f"{entry['file']}:{entry.get('mtime', 0)}")
|
||||
parts.append(f"random:{time.time_ns()}")
|
||||
else:
|
||||
for rel_path in _resolve_folder_selection(state, files):
|
||||
full_path = os.path.realpath(os.path.join(real_folder, rel_path))
|
||||
if not _is_within_directory(real_folder, full_path):
|
||||
parts.append(f"{rel_path}:outside")
|
||||
continue
|
||||
try:
|
||||
parts.append(f"{rel_path}:{os.stat(full_path).st_mtime_ns}")
|
||||
except OSError:
|
||||
parts.append(f"{rel_path}:missing")
|
||||
|
||||
return hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
class DumasH3PlanAttachSceneImagesNode:
|
||||
DESCRIPTION = (
|
||||
"Attach up to seven optional IMAGE sockets to one H3 Chain Plan scene "
|
||||
@@ -952,6 +1229,7 @@ class DumasCharacterHelperNode:
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DumasImageCompare": DumasImageCompareNode,
|
||||
"DumasSaveImage": DumasSaveImageNode,
|
||||
"DumasLoadImagesFolder": DumasLoadImagesFolderNode,
|
||||
"DumasH3PlanAttachSceneImages": DumasH3PlanAttachSceneImagesNode,
|
||||
"DumasH3PlanExtractSceneImages": DumasH3PlanExtractSceneImagesNode,
|
||||
"DumasCharacterHelper": DumasCharacterHelperNode,
|
||||
@@ -960,6 +1238,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DumasImageCompare": "Dumas Image Compare",
|
||||
"DumasSaveImage": "Save Image Dumas",
|
||||
"DumasLoadImagesFolder": "Load Images from Folder Dumas",
|
||||
"DumasH3PlanAttachSceneImages": "Dumas H3 Plan Attach Scene Images",
|
||||
"DumasH3PlanExtractSceneImages": "Dumas H3 Plan Extract Scene Images",
|
||||
"DumasCharacterHelper": "Dumas Character Helper",
|
||||
|
||||
Reference in New Issue
Block a user