Conor-Collins-ComfyUI-CoCoTools_IO/modules/image_loader.py

import os
import logging
import numpy as np
import torch
from PIL import Image, ImageOps
from typing import Tuple

# Import centralized logging setup
try:
    from ..utils.debug_utils import setup_logging
    setup_logging()
except ImportError:
    logging.basicConfig(level=logging.INFO)

logger = logging.getLogger(__name__)


class ImageLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_path": ("STRING", {
                    "default": "path/to/image.png",
                    "description": "Full path to the image file"
                }),
                "normalize": ("BOOLEAN", {
                    "default": True,
                    "description": "Normalize image values to the 0-1 range"
                })
            },
            "hidden": {"node_id": "UNIQUE_ID"}
        }

    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
    RETURN_NAMES = ("image", "mask", "metadata")
    FUNCTION = "load_regular_image"
    CATEGORY = "COCO Tools/Loaders"
    
    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")  # Always execute

    def load_regular_image(
        self, image_path: str, normalize: bool = True, node_id: str = None
    ) -> Tuple[torch.Tensor, torch.Tensor, str]:
        """
        Main function to load and process a regular image.
        Supports formats like PNG, JPG, and WebP.
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image path not found: {image_path}")

        try:
            with Image.open(image_path) as img:
                # Preserve original format before any transforms
                original_format = img.format

                # Apply EXIF orientation
                img = ImageOps.exif_transpose(img)

                # Detect bit depth from original mode BEFORE RGB conversion
                original_mode = img.mode
                info = self.detect_bit_depth(image_path, original_mode, original_format)
                bit_depth = info["bit_depth"]

                # Extract alpha channel before RGB conversion
                has_alpha = original_mode in ("RGBA", "LA", "PA")
                if has_alpha:
                    alpha_tensor = self.pil2tensor(img.split()[-1], 8).unsqueeze(0)

                # Convert to RGB and create tensor
                # convert("RGB") always produces 8-bit pixel data
                rgb_image = img.convert("RGB")
                rgb_tensor = self.pil2tensor(rgb_image, 8)

                # Default opaque alpha mask matching MASK format [B,H,W]
                if not has_alpha:
                    alpha_tensor = torch.ones(1, rgb_tensor.shape[1], rgb_tensor.shape[2])

                # Apply min-max range normalization if requested
                if normalize:
                    rgb_tensor = self.normalize_image(rgb_tensor)
                    alpha_tensor = self.normalize_image(alpha_tensor)

                # Prepare metadata
                metadata = {
                    "file_path": image_path,
                    "tensor_shape": tuple(rgb_tensor.shape),
                    "format": os.path.splitext(image_path)[1].lower(),
                    "bit_depth": bit_depth
                }

                return rgb_tensor, alpha_tensor, str(metadata)

        except Exception as e:
            logger.error(f"Error loading image {image_path}: {e}")
            raise ValueError(f"Error loading image {image_path}: {e}")

    @staticmethod
    def normalize_image(image: torch.Tensor) -> torch.Tensor:
        """
        Normalize a tensor to the 0-1 range via min-max stretching.
        Returns the image unchanged if all values are identical.
        """
        min_val, max_val = image.min(), image.max()
        if min_val == max_val:
            return image
        return (image - min_val) / (max_val - min_val)

    @staticmethod
    def detect_bit_depth(image_path: str, original_mode: str = None, original_format: str = None) -> dict:
        """
        Detect the bit depth of an image from its original mode before any conversion.
        """
        mode_to_bit_depth = {
            "1": 1, "L": 8, "P": 8, "RGB": 8, "RGBA": 8,
            "LA": 8, "PA": 8, "I;16": 16, "I": 32, "F": 32
        }

        if original_mode is not None:
            mode = original_mode
            fmt = original_format
        else:
            with Image.open(image_path) as img:
                mode = img.mode
                fmt = img.format

        bit_depth = mode_to_bit_depth.get(mode, 8)
        return {"bit_depth": bit_depth, "mode": mode, "format": fmt}

    @staticmethod
    def pil2tensor(image: Image.Image, bit_depth: int) -> torch.Tensor:
        """
        Convert a PIL Image to a PyTorch tensor, scaled to the 0-1 range.
        """
        image_np = np.array(image)
        if bit_depth == 8:
            image_tensor = torch.from_numpy(image_np.astype(np.float32) / 255.0)
        elif bit_depth == 16:
            image_tensor = torch.from_numpy(image_np.astype(np.float32) / 65535.0)
        elif bit_depth == 32:
            image_tensor = torch.from_numpy(image_np.astype(np.float32))
        else:
            logger.warning(f"Unsupported bit depth: {bit_depth}. Defaulting to 8-bit normalization.")
            image_tensor = torch.from_numpy(image_np.astype(np.float32) / 255.0)

        # Add a batch dimension if not present
        if len(image_tensor.shape) == 3:
            image_tensor = image_tensor.unsqueeze(0)

        return image_tensor