146 lines
5.0 KiB
Python
146 lines
5.0 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
import folder_paths
|
|
|
|
class IsulionLoadImagesNode:
|
|
"""
|
|
A flexible image loading node that preserves image characteristics
|
|
and prepares images for dynamic collage generation.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"default": "./input"}),
|
|
"target_row_height": ("INT", {"default": 300, "min": 100, "max": 1024, "step": 50}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "load_images"
|
|
CATEGORY = "Isulion/Image"
|
|
|
|
def load_images(self, directory, target_row_height=300):
|
|
"""
|
|
Load images from a directory with intelligent processing.
|
|
|
|
:param directory: Path to directory containing images
|
|
:param target_row_height: Target height for image rows
|
|
:return: Tensor of processed images
|
|
"""
|
|
# Resolve the full path
|
|
try:
|
|
full_directory = folder_paths.get_input_directory(directory)
|
|
except:
|
|
full_directory = os.path.abspath(os.path.expanduser(directory))
|
|
|
|
# Validate directory exists
|
|
if not os.path.isdir(full_directory):
|
|
raise ValueError(f"Directory does not exist: {full_directory}")
|
|
|
|
# Supported image extensions
|
|
image_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp']
|
|
|
|
# Get all image files in the directory
|
|
image_files = [
|
|
os.path.join(full_directory, f) for f in os.listdir(full_directory)
|
|
if os.path.isfile(os.path.join(full_directory, f)) and
|
|
os.path.splitext(f)[1].lower() in image_extensions
|
|
]
|
|
|
|
# Sort files for consistent order
|
|
image_files.sort()
|
|
|
|
# Check if any images were found
|
|
if not image_files:
|
|
raise ValueError(f"No images found in directory: {full_directory}")
|
|
|
|
# Process images
|
|
processed_images = []
|
|
|
|
for image_path in image_files:
|
|
# Open image
|
|
img = Image.open(image_path)
|
|
|
|
# Correct orientation
|
|
img = ImageOps.exif_transpose(img)
|
|
|
|
# Convert to RGB
|
|
img = img.convert('RGB')
|
|
|
|
# Calculate original aspect ratio
|
|
original_width, original_height = img.size
|
|
aspect_ratio = original_width / original_height
|
|
|
|
# Intelligent scaling to target row height
|
|
new_height = target_row_height
|
|
new_width = int(new_height * aspect_ratio)
|
|
|
|
# Resize maintaining aspect ratio
|
|
resized_img = img.resize((new_width, new_height), Image.LANCZOS)
|
|
|
|
# Convert to numpy array and normalize
|
|
img_array = np.array(resized_img).astype(np.float32) / 255.0
|
|
|
|
# Convert to tensor
|
|
img_tensor = torch.from_numpy(img_array)
|
|
|
|
processed_images.append(img_tensor)
|
|
|
|
# Convert to tensor without forcing same size
|
|
images_tensor = processed_images
|
|
|
|
return (images_tensor,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, directory, **kwargs):
|
|
"""
|
|
Check if the directory contents have changed.
|
|
"""
|
|
image_files = [
|
|
f for f in os.listdir(directory)
|
|
if os.path.isfile(os.path.join(directory, f)) and
|
|
os.path.splitext(f)[1].lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp']
|
|
]
|
|
return hash(tuple(sorted(image_files)))
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, directory, **kwargs):
|
|
"""
|
|
Validate the input directory.
|
|
|
|
:param directory: Directory to validate
|
|
:return: True if valid, error message if not
|
|
"""
|
|
try:
|
|
# Resolve the full path
|
|
try:
|
|
full_directory = folder_paths.get_input_directory(directory)
|
|
except:
|
|
full_directory = os.path.abspath(os.path.expanduser(directory))
|
|
|
|
# Check if directory exists
|
|
if not os.path.isdir(full_directory):
|
|
return f"Directory '{directory}' cannot be found."
|
|
|
|
# Supported image extensions
|
|
image_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp']
|
|
|
|
# Check if directory has any image files
|
|
image_files = [
|
|
f for f in os.listdir(full_directory)
|
|
if os.path.isfile(os.path.join(full_directory, f)) and
|
|
os.path.splitext(f)[1].lower() in image_extensions
|
|
]
|
|
|
|
if not image_files:
|
|
return f"No image files in directory '{directory}'."
|
|
|
|
return True
|
|
except Exception as e:
|
|
return f"Error validating directory: {str(e)}"
|