MdONeilsl-ComfyUI-MdColorMod/py/color_invert.py

import torch
from typing import Tuple, Optional

from .libs.color import srgb_to_linear_vec, linear_to_srgb_vec, rgb_to_hsv, hsv_to_rgb
from .libs.process import compute_chunk_size

# --------------------------------------------------------------------
# Internal helpers for channel handling and chunked processing
# --------------------------------------------------------------------

def _ensure_channels(t: torch.Tensor, target_channels: int) -> torch.Tensor:
    """Ensure tensor has exactly target_channels by adding/removing dummy dimensions.
       Used only for internal shape adjustments; for actual channel handling see below.
    """
    if t.ndim == 3 and t.shape[-1] == 1:
        # single channel -> expand to target? Not used here.
        pass
    return t


def _split_rgb_alpha(image: torch.Tensor):
    """Split image into RGB and alpha (if present)."""
    if image.shape[-1] == 3:
        return image, None
    elif image.shape[-1] == 4:
        return image[..., :3], image[..., 3:]
    else:
        raise ValueError(f"Unsupported number of channels: {image.shape[-1]}")


def _combine_rgb_alpha(rgb: torch.Tensor, alpha: Optional[torch.Tensor]) -> torch.Tensor:
    """Recombine RGB and alpha, or return RGB if no alpha."""
    if alpha is None:
        return rgb
    return torch.cat([rgb, alpha], dim=-1)


def _process_chunk(
    image_chunk: torch.Tensor,
    inv_type: str,
    mask_chunk: Optional[torch.Tensor]
) -> torch.Tensor:
    """
    Process a single chunk (batch of images) with the chosen inversion type.
    Handles RGB/RGBA, mask blending, and preserves alpha.
    """
    # Separate RGB and alpha
    rgb, alpha = _split_rgb_alpha(image_chunk)

    # Apply inversion to RGB channels only
    if inv_type == "Value":
        inverted_rgb = _exec_operator_color_value_invert_rgb(rgb)
    elif inv_type == "Linear":
        inverted_rgb = _exec_operator_color_linear_invert_rgb(rgb)
    else:  # "RGB"
        inverted_rgb = 1.0 - rgb

    # Apply mask blending if mask is provided
    if mask_chunk is not None:
        # Ensure mask has proper shape: (B, H, W) or (B, H, W, 1)
        if mask_chunk.ndim == 3:
            mask_chunk = mask_chunk.unsqueeze(-1)  # (B, H, W, 1)
        elif mask_chunk.ndim == 4 and mask_chunk.shape[-1] == 1:
            pass
        else:
            raise ValueError(f"Mask must have shape (B,H,W) or (B,H,W,1), got {mask_chunk.shape}")

        # Broadcast mask if needed (batch size 1 vs multiple)
        if mask_chunk.shape[0] == 1 and image_chunk.shape[0] > 1:
            mask_chunk = mask_chunk.expand(image_chunk.shape[0], -1, -1, -1)
        elif mask_chunk.shape[0] != image_chunk.shape[0]:
            raise ValueError(
                f"Mask batch size {mask_chunk.shape[0]} does not match image batch size {image_chunk.shape[0]}"
            )

        # Blend: output = image * (1 - mask) + inverted * mask
        rgb_blended = rgb * (1.0 - mask_chunk) + inverted_rgb * mask_chunk
        # Clamp to [0,1] for safety (though operations should keep within bounds)
        rgb_blended = torch.clamp(rgb_blended, 0.0, 1.0)
    else:
        rgb_blended = inverted_rgb

    # Recombine with original alpha
    return _combine_rgb_alpha(rgb_blended, alpha)


# --------------------------------------------------------------------
# Core inversion functions (unchanged logic, but work on RGB tensors)
# --------------------------------------------------------------------

def _exec_operator_color_linear_invert_rgb(rgb: torch.Tensor) -> torch.Tensor:
    """Linear inversion in RGB space (sRGB → linear → invert → sRGB)."""
    linear = srgb_to_linear_vec(rgb)
    inv_linear = 1.0 - linear
    return linear_to_srgb_vec(inv_linear)


def _exec_operator_color_value_invert_rgb(rgb: torch.Tensor) -> torch.Tensor:
    """Value (luminosity) inversion in HSV space."""
    h, s, v = rgb_to_hsv(rgb)
    v_inv = 1.0 - v
    inverted = hsv_to_rgb(h, s, v_inv)
    return torch.clamp(inverted, 0.0, 1.0)


# --------------------------------------------------------------------
# Public entry point
# --------------------------------------------------------------------

def exec_operator_color_invert(
    image: torch.Tensor,
    inv_type: str,
    mask: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor]:
    """
    Invert colours of an image batch, mimicking GIMP's Colors → Invert.

    Supports RGB and RGBA input/output; alpha channel is always passed through unchanged.
    If a mask is given, the inverted image is blended with the original using the mask
    as an alpha factor: output = image * (1 - mask) + inverted * mask. The mask must be
    a single-channel tensor with shape (B, H, W) or (B, H, W, 1). Its batch size must
    either match the image batch or be 1 (broadcast).

    Args:
        image: (B, H, W, C) tensor, C = 3 or 4, values in [0, 1].
        inv_type: One of "RGB", "Linear", or "Value".
        mask: Optional (B, H, W) or (B, H, W, 1) tensor, values in [0, 1].

    Returns:
        Tuple containing the inverted image tensor with same shape as input.

    Raises:
        ValueError: for invalid input shapes, types, or values.
    """
    # --- Input validation ---
    if not isinstance(image, torch.Tensor):
        raise TypeError("image must be a torch.Tensor")
    if image.ndim != 4:
        raise ValueError(f"image must be 4D (B,H,W,C), got {image.ndim}D")
    if image.shape[-1] not in (3, 4):
        raise ValueError(f"image must have 3 or 4 channels, got {image.shape[-1]}")
    if image.min() < 0 or image.max() > 1:
        raise ValueError("image values must be in [0, 1]")

    if inv_type not in ("RGB", "Linear", "Value"):
        raise ValueError(f"inv_type must be 'RGB', 'Linear', or 'Value', got {inv_type}")

    if mask is not None:
        if not isinstance(mask, torch.Tensor):
            raise TypeError("mask must be a torch.Tensor or None")
        if mask.ndim not in (3, 4) or (mask.ndim == 4 and mask.shape[-1] != 1):
            raise ValueError(f"mask must be (B,H,W) or (B,H,W,1), got {mask.shape}")
        if mask.min() < 0 or mask.max() > 1:
            raise ValueError("mask values must be in [0, 1]")

    # --- Early exit if mask is all zeros ---
    if mask is not None and torch.all(mask == 0):
        # No inversion applied – return original image (alpha preserved)
        return (image,)

    # --- Chunked processing ---
    batch_size = image.shape[0]
    if batch_size == 0:
        return (image,)

    # Compute optimal chunk size using the provided helper
    chunk_size = compute_chunk_size(image, mask)
    if chunk_size <= 0:
        chunk_size = batch_size  # fallback

    # Process in chunks
    chunks = []
    for start in range(0, batch_size, chunk_size):
        end = min(start + chunk_size, batch_size)
        img_chunk = image[start:end]
        mask_chunk = mask[start:end] if mask is not None else None
        processed = _process_chunk(img_chunk, inv_type, mask_chunk)
        chunks.append(processed)

    # Concatenate chunks along batch dimension
    result = torch.cat(chunks, dim=0)
    return (result,)