59 lines
1.5 KiB
Python
59 lines
1.5 KiB
Python
"""
|
|
Image loading utilities
|
|
"""
|
|
|
|
import os
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Optional
|
|
|
|
try:
|
|
from folder_paths import get_input_directory
|
|
except ImportError:
|
|
# Fallback for testing
|
|
def get_input_directory():
|
|
return "input"
|
|
|
|
|
|
def load_image_from_disk(
|
|
filename: str,
|
|
subfolder: str = ""
|
|
) -> Optional[np.ndarray]:
|
|
"""
|
|
Load image from ComfyUI input directory.
|
|
|
|
Args:
|
|
filename: Image filename
|
|
subfolder: Optional subfolder within input directory
|
|
|
|
Returns:
|
|
Image array (H, W, 3) in range [0, 1] as float32, or None if failed
|
|
"""
|
|
if not filename:
|
|
return None
|
|
|
|
try:
|
|
input_dir = get_input_directory()
|
|
if subfolder:
|
|
image_path = os.path.join(input_dir, subfolder, filename)
|
|
else:
|
|
image_path = os.path.join(input_dir, filename)
|
|
|
|
if not os.path.exists(image_path):
|
|
print(f"[RotoMaskNode] Image not found: {image_path}")
|
|
return None
|
|
|
|
# Load image using PIL
|
|
img = Image.open(image_path)
|
|
# Convert to RGB if needed
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
|
|
# Convert to numpy array and normalize to [0, 1]
|
|
img_array = np.array(img, dtype=np.float32) / 255.0
|
|
|
|
return img_array
|
|
except Exception as e:
|
|
print(f"[RotoMaskNode] Error loading image {filename}: {e}")
|
|
return None
|