130 lines
5.2 KiB
Python
130 lines
5.2 KiB
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps
|
|
|
|
class LoadLoopImagesFromDir:
|
|
last_index = 0 # Class attribute to track the last accessed image
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
The function `INPUT_TYPES` defines required and optional input types with default values and
|
|
labels for a specified class.
|
|
|
|
:param cls: The `INPUT_TYPES` function defines input parameters for a class. The parameters are
|
|
categorized as "required" and "optional" inputs
|
|
:return: A dictionary containing information about input types for a class. The dictionary has
|
|
two keys: "required" and "optional", each containing a dictionary of input types with their
|
|
corresponding data types and optional parameters such as default values and labels.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"directory": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"load_always": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "STRING")
|
|
RETURN_NAMES = ("IMAGE", "MASK", "FILE PATH")
|
|
OUTPUT_IS_LIST = (False, False, False) # Only return one image at a time
|
|
|
|
FUNCTION = "load_images"
|
|
|
|
CATEGORY = "image"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
"""
|
|
The function `IS_CHANGED` increments `last_index` every time it is called, looping when needed,
|
|
and returns a hash based on the index and input arguments.
|
|
|
|
:param cls: The `cls` parameter in the given function represents the class itself. It is used to
|
|
access class variables and methods within the function
|
|
:return: The function `IS_CHANGED` returns the hash value of a tuple containing `cls.last_index`
|
|
and a frozenset of the keyword arguments `kwargs.items()`.
|
|
"""
|
|
if kwargs.get("load_always", False):
|
|
return float("NaN")
|
|
|
|
directory = kwargs.get("directory", "")
|
|
if not os.path.isdir(directory):
|
|
return hash(frozenset(kwargs.items()))
|
|
|
|
valid_extensions = {'.jpg', '.jpeg', '.png', '.webp'}
|
|
dir_files = sorted(
|
|
f for f in os.listdir(directory)
|
|
if any(f.lower().endswith(ext) for ext in valid_extensions)
|
|
)
|
|
|
|
if dir_files:
|
|
cls.last_index = (cls.last_index + 1) % len(dir_files)
|
|
|
|
return hash((cls.last_index, frozenset(kwargs.items())))
|
|
|
|
def load_images(cls, directory: str, load_always=False):
|
|
"""
|
|
The function `load_images` loads image files from a specified directory, processes them, and
|
|
returns the
|
|
|
|
image, mask (if alpha channel exists), and the image path.
|
|
|
|
:param directory: The `directory` parameter in the `load_images` method is a string that
|
|
represents the path to the directory containing the image files that you want to load. This
|
|
method loads images from the specified directory based on the valid image file extensions (.jpg,
|
|
.jpeg, .png, .webp)
|
|
:type directory: str
|
|
:param load_always: The `load_always` parameter in the `load_images` method is a boolean flag
|
|
that determines whether to always load images from the specified directory or not. By default,
|
|
it is set to `False`, which means that images will be loaded cyclically from the directory based
|
|
on the `, defaults to False (optional)
|
|
:return: The `load_images` method returns three values: `image`, `mask`, and `image_path`.
|
|
"""
|
|
if not os.path.isdir(directory):
|
|
raise FileNotFoundError(f"Directory '{directory}' not found.")
|
|
|
|
# Get all valid image files
|
|
valid_extensions = {'.jpg', '.jpeg', '.png', '.webp'}
|
|
dir_files = sorted(
|
|
f for f in os.listdir(directory)
|
|
if any(f.lower().endswith(ext) for ext in valid_extensions)
|
|
)
|
|
|
|
if not dir_files:
|
|
raise FileNotFoundError(f"No image files found in directory '{directory}'.")
|
|
|
|
# Cycle through images using `last_index`
|
|
image_count = len(dir_files)
|
|
image_path = os.path.join(directory, dir_files[cls.last_index])
|
|
|
|
# Update index for the next call (loop back if at the end)
|
|
cls.last_index = (cls.last_index + 1) % image_count
|
|
|
|
# Load image
|
|
img = Image.open(image_path)
|
|
img = ImageOps.exif_transpose(img)
|
|
image = img.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,] # Add batch dimension
|
|
|
|
# Load mask (if alpha channel exists)
|
|
if 'A' in img.getbands():
|
|
mask = np.array(img.getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1.0 - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
|
|
return image, mask, str(image_path)
|
|
|
|
# Dictionary to map node classes to their names
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadLoopImagesFromDir": LoadLoopImagesFromDir,
|
|
}
|
|
|
|
# Dictionary to map node names to their friendly display names
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadLoopImagesFromDir": "Load Loop Images From Dir",
|
|
}
|