RemoveBackgroundAndResizeNode.py

"""
RemoveBackgroundAndResizeNode
============================
Single-file ComfyUI node for background removal with integrated resize.

Based on ``TransparencyBackgroundRemover`` and ``ScalingMixin`` from the
ComfyUI-TransparencyBackgroundRemover package.  Adapted for standalone use
with minimal dependencies — PIL only for background removal, no hard
OpenCV requirement.

Author:  Gero Doll / Limbicnation
License: Apache-2.0
"""

from __future__ import annotations

import logging
from typing import Tuple

import numpy as np
import torch
from PIL import Image

try:
    import structlog
    log = structlog.get_logger(__name__)
except ImportError:
    log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# ScalingMixin — shared resize logic (see scaling.py)
# ---------------------------------------------------------------------------

try:
    from .scaling import ScalingMixin  # noqa: F401 — re-exported for convenience
except ImportError:
    from scaling import ScalingMixin


# ---------------------------------------------------------------------------
# PIL-only background removal
# ---------------------------------------------------------------------------

def _kmeans_colors(
    arr: np.ndarray, n_clusters: int, tol: float = 0.01, max_iter: int = 20,
    max_samples: int = 50000,
) -> np.ndarray:
    """
    Simple k-means on pixel array [H*W, channels].
    Returns ``(n_clusters, channels)`` cluster centres.
    Fallback when scikit-learn is unavailable.

    Parameters
    ----------
    arr : np.ndarray
        Input image array (H, W, C).
    n_clusters : int
        Number of k-means clusters.
    tol : float
        Convergence tolerance for centre shift.
    max_iter : int
        Maximum k-means iterations.
    max_samples : int
        Cap on number of pixels fed to k-means to avoid OOM on large images.
        Subsampled uniformly when pixel count exceeds this value.
    """
    pixels = arr.reshape(-1, arr.shape[-1]).astype(np.float32)

    # Subsample to avoid OOM on large images (e.g. 4K → 8.3M pixels)
    if len(pixels) > max_samples:
        rng = np.random.RandomState(42)
        idx = rng.choice(len(pixels), size=max_samples, replace=False)
        pixels = pixels[idx]

    indices = np.random.RandomState(42).choice(
        len(pixels), size=n_clusters, replace=False
    )
    centres = pixels[indices].copy()

    for _ in range(max_iter):
        dists = np.linalg.norm(pixels[:, None] - centres[None], axis=-1)
        labels = dists.argmin(axis=1)
        new_centres = np.stack([pixels[labels == i].mean(axis=0) for i in range(n_clusters)])
        if np.abs(new_centres - centres).max() < tol:
            break
        centres = new_centres

    return centres


def _dominant_bg(arr: np.ndarray, samples: int = 5000) -> np.ndarray:
    """Sample corners + border to estimate background colour."""
    h, w = arr.shape[:2]
    mask = np.zeros((h, w), dtype=bool)
    bw = max(1, min(h, w) // 10)

    # top / bottom strips and left / right strips (intersection gives the four corners)
    mask[:bw, :] = True
    mask[-bw:, :] = True
    mask[:, :bw] = True
    mask[:, -bw:] = True

    bg_pixels = arr[mask]
    if len(bg_pixels) > samples:
        idx = np.random.RandomState(42).choice(len(bg_pixels), samples, replace=False)
        bg_pixels = bg_pixels[idx]

    return bg_pixels.mean(axis=0).astype(np.uint8)


def remove_background_pil(
    image: Image.Image,
    tolerance: int = 30,
    color_clusters: int = 8,
    foreground_bias: float = 0.7,
    binary_threshold: int = 128,
    dither_handling: bool = True,
) -> Image.Image:
    """
    Remove background from a PIL Image using colour clustering.

    Pipeline
    --------
    1. Flatten to ``[H*W, channels]`` and run k-means to find dominant colours.
    2. Identify background as the colour most similar to image edges / corners.
    3. Compute per-pixel distance to foreground centres; keep pixels that are
       closer than ``tolerance`` (adjusted by ``foreground_bias``).
    4. Apply simple edge refinement via binary thresholding.
    5. Return ``RGBA`` PIL Image.

    Parameters
    ----------
    image : PIL.Image
        Input image (RGB or RGBA; RGBA ``A`` is ignored).
    tolerance : int
        Colour distance threshold (0-255). Higher = more pixels accepted as foreground.
    color_clusters : int
        K-means centres used to identify distinct colours.
    foreground_bias : float
        Bias towards keeping pixels (0 = discard more, 1 = keep more).
    binary_threshold : int
        Threshold for alpha mask binarisation (0-255).
    dither_handling : bool
        When True, apply a 1-pixel border erosion to reduce dither artefacts.

    Returns
    -------
    PIL.Image
        RGBA image with transparent background.
    """
    rgb = image.convert("RGB")
    arr = np.array(rgb, dtype=np.float32)  # [H, W, 3]

    # 1. k-means colour clustering
    if color_clusters < 2:
        color_clusters = 2
    centres = _kmeans_colors(arr, color_clusters)

    # 2. background colour from image perimeter
    bg_colour = _dominant_bg(arr)

    # 3. distance of each pixel to background vs foreground centres
    h, w = arr.shape[:2]
    flat = arr.reshape(-1, 3)

    bg_dist = np.linalg.norm(flat - bg_colour.astype(np.float32), axis=1)
    fg_dist = np.linalg.norm(
        flat[:, None] - centres.astype(np.float32), axis=-1
    ).min(axis=1)

    # effective tolerance — foreground_bias shifts the acceptance boundary
    effective_tol = tolerance * (1.5 - foreground_bias)
    mask = (fg_dist + effective_tol) < bg_dist

    alpha = mask.reshape(h, w).astype(np.uint8) * 255

    # 4. edge refinement
    if dither_handling:
        from PIL import ImageFilter

        alpha_pil = Image.fromarray(alpha, mode="L")
        alpha_pil = alpha_pil.filter(ImageFilter.MinFilter(3))
        alpha = np.array(alpha_pil)

    # 5. binarise
    alpha = (alpha > binary_threshold).astype(np.uint8) * 255

    # 6. assemble RGBA
    result = np.dstack([arr.astype(np.uint8), alpha])
    return Image.fromarray(result, mode="RGBA")


# ---------------------------------------------------------------------------
# ComfyUI node
# ---------------------------------------------------------------------------

class RemoveBackgroundAndResizeNode(ScalingMixin):
    """
    Combines background removal and resize in a single node.

    ``INPUT_TYPES`` exposes tolerance, edge sensitivity, colour-cluster, and
    output-size controls so the node can be dropped into any workflow without
    separate background-removal and resize chains.

    Outputs
    -------
    IMAGE : RGBA tensor (N, H, W, 4), values 0-1
    MASK  : greyscale alpha tensor (N, H, W),    values 0-1
    """

    # ------------------------------------------------------------------
    # ComfyUI contract
    # ------------------------------------------------------------------
    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("image", "mask")
    FUNCTION = "process"
    CATEGORY = "image/processing"
    OUTPUT_NODE = True

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE", {"tooltip": "Input image tensor from upstream node"}),
                "tolerance": (
                    "INT",
                    {
                        "default": 30,
                        "min": 0,
                        "max": 255,
                        "step": 1,
                        "display": "number",
                        "tooltip": (
                            "Colour distance threshold. Higher = more aggressive "
                            "foreground extraction (0-255)."
                        ),
                    },
                ),
                "foreground_bias": (
                    "FLOAT",
                    {
                        "default": 0.7,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "display": "slider",
                        "tooltip": (
                            "Bias toward keeping pixels as foreground. "
                            "Higher = more pixels preserved."
                        ),
                    },
                ),
                "color_clusters": (
                    "INT",
                    {
                        "default": 8,
                        "min": 2,
                        "max": 20,
                        "step": 1,
                        "display": "number",
                        "tooltip": "Number of k-means centres for colour segmentation.",
                    },
                ),
                "output_size": (
                    [
                        "ORIGINAL",
                        "512x512",
                        "768x768",
                        "1024x1024",
                        "1280x720",
                        "1920x1080",
                    ],
                    {
                        "default": "ORIGINAL",
                        "tooltip": (
                            "Resize output to this size. "
                            "ORIGINAL keeps source dimensions."
                        ),
                    },
                ),
                "scaling_method": (
                    ["NEAREST", "BILINEAR", "BICUBIC", "LANCZOS"],
                    {
                        "default": "LANCZOS",
                        "tooltip": (
                            "Interpolation method. "
                            "LANCZOS = best quality; "
                            "NEAREST = pixel-perfect for pixel art."
                        ),
                    },
                ),
            },
            "optional": {
                "dither_handling": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Apply edge erosion to suppress dither artefacts.",
                    },
                ),
                "binary_threshold": (
                    "INT",
                    {
                        "default": 128,
                        "min": 0,
                        "max": 255,
                        "step": 1,
                        "display": "number",
                        "tooltip": "Alpha mask binarisation cutoff (0-255).",
                    },
                ),
            },
        }

    # ------------------------------------------------------------------
    # processing
    # ------------------------------------------------------------------

    @torch.no_grad()
    def process(
        self,
        image: torch.Tensor,
        tolerance: int,
        foreground_bias: float,
        color_clusters: int,
        output_size: str,
        scaling_method: str,
        dither_handling: bool = True,
        binary_threshold: int = 128,
    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
        """
        ``image`` — ComfyUI IMAGE tensor, shape (N, H, W, C), float32 0-1.
        Returns (IMAGE tensor, MASK tensor).
        """
        # ------------------------------------------------------------------
        # validation
        # ------------------------------------------------------------------
        if image is None or image.numel() == 0:
            raise ValueError("RemoveBackgroundAndResizeNode: received empty image tensor.")
        if image.ndim != 4:
            raise ValueError(
                f"RemoveBackgroundAndResizeNode: expected 4-D IMAGE tensor "
                f"(N,H,W,C), got {image.ndim}-D."
            )
        if image.shape[-1] not in (3, 4):
            raise ValueError(
                f"RemoveBackgroundAndResizeNode: IMAGE must be RGB or RGBA; "
                f"got channels={image.shape[-1]}."
            )

        # ------------------------------------------------------------------
        # per-batch processing
        # ------------------------------------------------------------------
        rgba_tensors: list[torch.Tensor] = []
        mask_tensors: list[torch.Tensor] = []

        for i in range(image.shape[0]):
            img_np = (image[i].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)

            # handle RGB input
            if img_np.shape[-1] == 3:
                pil_in = Image.fromarray(img_np, mode="RGB")
            else:
                pil_in = Image.fromarray(img_np[:, :, :3], mode="RGB")

            # ---- remove background ----
            pil_rgba = remove_background_pil(
                pil_in,
                tolerance=tolerance,
                foreground_bias=foreground_bias,
                color_clusters=color_clusters,
                binary_threshold=binary_threshold,
                dither_handling=dither_handling,
            )

            # ---- resize ----
            pil_rgba = self.resize_image(pil_rgba, output_size, scaling_method)

            # ---- back to tensor ----
            arr = np.array(pil_rgba, dtype=np.float32) / 255.0
            rgba_tensors.append(torch.from_numpy(arr))

            # extract alpha channel as mask
            mask_np = (np.array(pil_rgba.split()[3], dtype=np.float32) / 255.0)
            mask_tensors.append(torch.from_numpy(mask_np))

        # ------------------------------------------------------------------
        # stack and return in ComfyUI tensor format
        # ------------------------------------------------------------------
        stacked_image: torch.Tensor = torch.stack(rgba_tensors).to(dtype=torch.float32)
        stacked_mask: torch.Tensor = torch.stack(mask_tensors).to(dtype=torch.float32)

        # ComfyUI MASK convention: (N, H, W)
        if stacked_mask.ndim == 4 and stacked_mask.shape[-1] == 1:
            stacked_mask = stacked_mask.squeeze(-1)

        log.debug(
            "RemoveBackgroundAndResizeNode.done",
            batch=image.shape[0],
            output_size=output_size,
            scaling_method=scaling_method,
        )

        return (stacked_image, stacked_mask)


# ComfyUI registration dict
NODE_CLASS_MAPPINGS = {"RemoveBackgroundAndResize": RemoveBackgroundAndResizeNode}
NODE_DISPLAY_NAME_MAPPINGS = {
    "RemoveBackgroundAndResize": "Remove Background & Resize"
}