109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
class FolderImageLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"default": ""}),
|
|
"target_width": ("INT", {"default": 512, "min": 0, "max": 8192, "step": 1, "tooltip": "0 = use first image width"}),
|
|
"target_height": ("INT", {"default": 512, "min": 0, "max": 8192, "step": 1, "tooltip": "0 = use first image height"}),
|
|
"image_load_cap": ("INT", {"default": 0, "min": 0, "max": 9999, "step": 1, "tooltip": "0 = unlimited"}),
|
|
"start_index": ("INT", {"default": 0, "min": 0, "max": 9999, "step": 1, "tooltip": "Skip first N images"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "STRING")
|
|
RETURN_NAMES = ("images", "masks", "filenames")
|
|
OUTPUT_IS_LIST = (True, True, True)
|
|
FUNCTION = "load_images"
|
|
CATEGORY = "AnotherUtils"
|
|
|
|
def load_images(self, directory, target_width, target_height, image_load_cap, start_index):
|
|
if not os.path.isdir(directory):
|
|
raise FileNotFoundError(f"Directory not found: {directory}")
|
|
|
|
valid_ext = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp')
|
|
files = [f for f in os.listdir(directory) if f.lower().endswith(valid_ext)]
|
|
files.sort()
|
|
|
|
# Apply start index
|
|
if start_index > 0:
|
|
files = files[start_index:]
|
|
|
|
# Apply Cap
|
|
if image_load_cap > 0:
|
|
files = files[:image_load_cap]
|
|
|
|
if not files:
|
|
raise FileNotFoundError(f"No valid images found in: {directory}")
|
|
|
|
images = []
|
|
masks = []
|
|
filenames = []
|
|
|
|
# Determine strict target size
|
|
if target_width == 0 or target_height == 0:
|
|
# Fallback to first image size for 0 dims
|
|
try:
|
|
temp = Image.open(os.path.join(directory, files[0]))
|
|
if target_width == 0:
|
|
target_width = temp.size[0]
|
|
if target_height == 0:
|
|
target_height = temp.size[1]
|
|
temp.close()
|
|
except Exception:
|
|
target_width = 512
|
|
target_height = 512
|
|
|
|
final_size = (target_width, target_height)
|
|
print(f"[FolderImageLoader] Loading {len(files)} images. Canvas Size: {final_size}")
|
|
|
|
for f in files:
|
|
path = os.path.join(directory, f)
|
|
try:
|
|
img = Image.open(path)
|
|
|
|
# Handle Animation
|
|
if getattr(img, 'is_animated', False):
|
|
img.seek(0)
|
|
|
|
# Logic: Canvas Center with White Background
|
|
# Create white background canvas
|
|
canvas = Image.new("RGBA", final_size, (255, 255, 255, 255))
|
|
|
|
# Calculate center position
|
|
paste_x = (final_size[0] - img.size[0]) // 2
|
|
paste_y = (final_size[1] - img.size[1]) // 2
|
|
|
|
# Paste using alpha mask
|
|
canvas.paste(img, (paste_x, paste_y), mask=img.split()[3] if img.mode == 'RGBA' else None)
|
|
img = canvas
|
|
|
|
# Process Image
|
|
# Normalize to 0-1 float
|
|
# (H, W, 4)
|
|
i = np.array(img).astype(np.float32) / 255.0
|
|
|
|
mask = i[:, :, 3]
|
|
image_rgb = i[:, :, :3]
|
|
|
|
# Add [1, ...] dimension for ComfyUI Image format if returning list
|
|
# Each item in list must be [1, H, W, C]
|
|
images.append(torch.from_numpy(image_rgb).unsqueeze(0))
|
|
masks.append(torch.from_numpy(mask)) # Shape [H, W] for ComfyUI compatibility
|
|
|
|
filenames.append(f)
|
|
|
|
except Exception as e:
|
|
print(f"[FolderImageLoader] Error loading {f}: {e}")
|
|
|
|
if not images:
|
|
raise ValueError("Failed to process any images.")
|
|
|
|
print(f"[FolderImageLoader] Created list of {len(images)} images")
|
|
return (images, masks, filenames)
|