Initial commit

image/__init__.py

@@ -0,0 +1,3 @@
+"""
+Image processing modules
+"""

image/loader.py

@@ -0,0 +1,58 @@
+"""
+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

image/processor.py

@@ -0,0 +1,66 @@
+"""
+Image processing functions
+"""
+
+import numpy as np
+import cv2
+from typing import List
+
+
+def resize_mask_to_image(
+    mask: np.ndarray,
+    target_height: int,
+    target_width: int
+) -> np.ndarray:
+    """
+    Resize mask to match image dimensions.
+    
+    Args:
+        mask: Mask array (H, W) in range [0, 1]
+        target_height: Target height
+        target_width: Target width
+        
+    Returns:
+        Resized mask array (target_height, target_width) in range [0, 1]
+    """
+    if mask.shape[0] == target_height and mask.shape[1] == target_width:
+        return mask
+    
+    # Use cv2.resize with INTER_LINEAR for smooth resizing
+    mask_resized = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
+    # Ensure values stay in [0, 1] range
+    mask_resized = np.clip(mask_resized, 0.0, 1.0)
+    
+    return mask_resized
+
+
+def apply_mask_overlay(
+    image: np.ndarray,
+    mask: np.ndarray,
+    background_color: List[float]
+) -> np.ndarray:
+    """
+    Apply mask overlay to image with background color.
+    Where mask is 1.0, use background_color; where mask is 0.0, use original image.
+    
+    Args:
+        image: Image array (H, W, 3) in range [0, 1]
+        mask: Mask array (H, W) in range [0, 1]
+        background_color: RGB color [r, g, b] in range [0, 1]
+        
+    Returns:
+        Image with mask overlay applied (H, W, 3) in range [0, 1]
+    """
+    # Ensure mask is 3D for broadcasting: (H, W) -> (H, W, 1)
+    mask_3d = np.expand_dims(mask, axis=-1)
+    
+    # Expand background_color to match image shape: (3,) -> (1, 1, 3)
+    bg_color_array = np.array(background_color, dtype=np.float32).reshape(1, 1, 3)
+    
+    # Apply overlay: image * (1 - mask) + background_color * mask
+    masked_image = image * (1.0 - mask_3d) + bg_color_array * mask_3d
+    
+    # Ensure values stay in [0, 1] range
+    masked_image = np.clip(masked_image, 0.0, 1.0)
+    
+    return masked_image

image/utils.py

@@ -0,0 +1,24 @@
+"""
+Image utility functions
+"""
+
+import torch
+
+
+def create_empty_mask(
+    batch_size: int,
+    height: int,
+    width: int
+) -> torch.Tensor:
+    """
+    Create an empty mask tensor.
+    
+    Args:
+        batch_size: Number of images in batch
+        height: Mask height
+        width: Mask width
+        
+    Returns:
+        Zero-filled mask tensor as RGB (B, H, W, 3)
+    """
+    return torch.zeros((batch_size, height, width, 3), dtype=torch.float32)