diffmat.py

#!/usr/bin/env python3
"""
Remove backgrounds from AI-generated images via multi-background pixel comparison.

AI image generators (e.g. Nano Banana, Midjourney) cannot output true transparency.
This tool solves that by comparing 2-5 renders of the same subject on different solid
background colours. Pixels that match their assigned background across ALL images become
transparent; pixels that are consistent across images become the subject; transitional
pixels get partial alpha for smooth edges.

Three matting algorithms are available (see --method):
  simple         — Binary classification. Fast, good for crisp subjects.
  variance       — Statistical variance model. Smoother anti-aliased edges.
  decomposition  — Linear unmixing (pixel = a*subject + (1-a)*bg). Best for
                   semi-transparent edges, drop shadows, and glass effects.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np
from PIL import Image


# =============================================================================
# Constants
# =============================================================================

MIN_IMAGES = 2
MAX_IMAGES = 5

DEFAULT_BG_COLORS: List[Tuple[int, int, int]] = [
    (255, 255, 255),  # white
    (0, 0, 0),        # black
    (255, 0, 0),      # red
    (0, 255, 0),      # green
    (0, 0, 255),      # blue
]

_COLOR_KEYWORDS: dict[Tuple[int, int, int], List[str]] = {
    (255, 255, 255): ["white", "whit"],
    (0, 0, 0):       ["black", "blac"],
    (255, 0, 0):     ["red"],
    (0, 255, 0):     ["green", "gree"],
    (0, 0, 255):     ["blue", "blu"],
}

_OUTPUT_EXCLUDE_PATTERNS = ("output.png", "output_", "_transparent.png", "test_")


# =============================================================================
# Image I/O
# =============================================================================

def load_images(paths: List[Path]) -> List[np.ndarray]:
    """Load image files as RGB numpy arrays of shape (H, W, 3)."""
    result = []
    for p in paths:
        result.append(np.array(Image.open(p).convert("RGB")))
    return result


def verify_image_sizes(images: List[np.ndarray]) -> Tuple[int, int]:
    """Return (H, W) after confirming all images share the same dimensions."""
    if not images:
        raise ValueError("No images provided")
    shape = images[0].shape[:2]
    for idx, img in enumerate(images[1:], 2):
        if img.shape[:2] != shape:
            raise ValueError(
                f"Image {idx} is {img.shape[:2]} but image 1 is {shape}; "
                "all images must have identical dimensions."
            )
    return shape


def find_images_in_folder(folder: Path) -> List[Path]:
    """Return sorted list of image paths in *folder*, excluding prior outputs."""
    extensions = [".png", ".jpg", ".jpeg", ".PNG", ".JPG", ".JPEG"]
    files: List[Path] = []
    for ext in extensions:
        files.extend(folder.glob(f"*{ext}"))

    def _is_output(p: Path) -> bool:
        name = p.name.lower()
        return (
            name == "output.png"
            or name.startswith("output_")
            or name.endswith("_transparent.png")
            or any(name.startswith(pat) for pat in _OUTPUT_EXCLUDE_PATTERNS if pat.endswith("_"))
        )

    return sorted(p for p in files if not _is_output(p))


# =============================================================================
# Colour helpers
# =============================================================================

def color_distance(a: Tuple[int, int, int], b: Tuple[int, int, int]) -> float:
    """Euclidean distance between two RGB triples (0-255)."""
    return float(np.sqrt(sum((float(x) - float(y)) ** 2 for x, y in zip(a, b))))


# =============================================================================
# Algorithm 1 — Simple (binary classification)
# =============================================================================
# Decision per pixel:
#   1. Every image shows its own background colour  -> transparent
#   2. All images show a similar colour              -> opaque (subject)
#   3. Otherwise                                     -> edge
#
# Edge alpha:
#   hard     -> majority vote (>= N/2 match background => 0, else 255)
#   smooth   -> average of normalised distance-to-background, clamped 0-255
#
# Pros:  Fast, predictable, great for crisp subjects with hard edges.
# Cons:  Can produce jagged silhouettes; struggles with semi-transparent
#        regions (glass, smoke, soft drop-shadows) because the binary
#        similar / not-similar split is too coarse.

def _method_simple(
    images: list[np.ndarray],
    bg_colors: list[Tuple[int, int, int]],
    tol: float,
    sim_tol: float,
    ref_idx: int,
    aa: bool,
) -> np.ndarray:
    """Binary classification background removal."""
    n = len(images)
    h, w = verify_image_sizes(images)
    base = images[ref_idx].copy()
    out = np.zeros((h, w, 4), dtype=np.uint8)
    out[:, :, :3] = base
    alpha = np.zeros((h, w), dtype=np.float32)

    for y in range(h):
        for x in range(w):
            pxs = [tuple(img[y, x]) for img in images]

            # --- transparent? ---
            if all(color_distance(pxs[i], bg_colors[i]) <= tol for i in range(n)):
                continue  # alpha stays 0

            # --- opaque subject? ---
            if all(color_distance(pxs[0], pxs[i]) <= sim_tol for i in range(1, n)):
                alpha[y, x] = 255.0
                continue

            # --- edge ---
            if not aa:
                bg_hits = sum(1 for i in range(n) if color_distance(pxs[i], bg_colors[i]) <= tol)
                alpha[y, x] = 0.0 if bg_hits >= n / 2 else 255.0
            else:
                strength = sum(
                    min(1.0, color_distance(pxs[i], bg_colors[i]) / tol)
                    for i in range(n)
                ) / n
                alpha[y, x] = np.clip(255.0 * strength, 0.0, 255.0)

    out[:, :, 3] = alpha.astype(np.uint8)
    return out


# =============================================================================
# Algorithm 2 — Variance-based matting
# =============================================================================
# For each pixel we compute the variance across all N images:
#   low variance + far from every background  -> opaque
#   low variance + near every background      -> transparent
#   medium variance                           -> partial alpha
#
# Instead of "similar or not" we use the actual spread of values.
# Alpha is a weighted blend of (1 - normalised_variance) and
# normalised_distance_to_background.
#
# Pros:  Smoother anti-aliased edges than *simple* because the transition
#        from opaque to transparent is driven by a continuous statistic
#        (variance) rather than a threshold.
# Cons:  Can be too soft on very fine detail (hair, thin lines);
#        variance alone doesn't tell *why* pixels differ (shadow vs edge).

def _method_variance(
    images: list[np.ndarray],
    bg_colors: list[Tuple[int, int, int]],
    tol: float,
    sim_tol: float,
    ref_idx: int,
    aa: bool,
) -> np.ndarray:
    """Variance-based matting."""
    n = len(images)
    h, w = verify_image_sizes(images)
    base = images[ref_idx].copy()

    stacked = np.stack(images, axis=0).astype(np.float64)          # (n,h,w,3)
    bg_arr   = np.array(bg_colors, dtype=np.float64).reshape(n,1,1,3)

    var_rgb  = np.var(stacked, axis=0)                             # (h,w,3)
    var_mean = np.mean(var_rgb, axis=2)                            # (h,w)
    dist_bg  = np.sqrt(np.sum((stacked - bg_arr) ** 2, axis=3))    # (n,h,w)
    all_bg   = np.all(dist_bg <= tol, axis=0)                      # (h,w)
    pix_mean = np.mean(stacked, axis=0)                            # (h,w,3)
    max_dev  = np.max(
        np.sqrt(np.sum((stacked - pix_mean[None, :, :, :]) ** 2, axis=3)), axis=0
    )

    out   = np.zeros((h, w, 4), dtype=np.uint8)
    out[:, :, :3] = base
    alpha = np.zeros((h, w), dtype=np.float32)

    v_sub = (sim_tol / 2) ** 2
    v_bg  = (tol / 2) ** 2

    for y in range(h):
        for x in range(w):
            if all_bg[y, x]:
                continue

            v = var_mean[y, x]
            d = max_dev[y, x]

            if v <= v_sub and d <= sim_tol:
                alpha[y, x] = 255.0
            elif not aa:
                hits = sum(1 for i in range(n) if dist_bg[i, y, x] <= tol)
                alpha[y, x] = 0.0 if hits >= n / 2 else 255.0
            else:
                nv = min(1.0, v / v_bg) if v_bg > 0 else 1.0
                df = min(1.0, np.mean([dist_bg[i, y, x] for i in range(n)]) / tol) if tol > 0 else 1.0
                alpha[y, x] = np.clip(255.0 * ((1 - nv) * 0.5 + df * 0.5), 0.0, 255.0)

    out[:, :, 3] = alpha.astype(np.uint8)
    return out


# =============================================================================
# Algorithm 3 — Decomposition (linear unmixing)
# =============================================================================
# Physical model per pixel:
#     I_i = a * S + (1 - a) * B_i          for each image i
#   I_i = observed pixel colour in image i
#   S   = true subject colour (unknown)
#   B_i = known background colour of image i
#   a   = alpha (opacity) we want to recover
#
# We solve iteratively:
#   1. Seed S with the pixel furthest from its background.
#   2. Given a, solve for S from each image, weighted by how far that
#      pixel is from its background (far = more reliable).
#   3. Given S, solve for a via least-squares projection.
#   4. Repeat steps 2-3 a few times until convergence.
#
# Pros:  Best quality for semi-transparent edges, drop shadows, glass,
#        and smoke. Recovers the *true subject colour* by subtracting
#        the estimated background contribution from each image.
# Cons:  Slowest (iterative per pixel). Can over-estimate alpha on noisy
#        / compressed images where JPEG artefacts mimic background bleed.

def _method_decomposition(
    images: list[np.ndarray],
    bg_colors: list[Tuple[int, int, int]],
    tol: float,
    sim_tol: float,
    ref_idx: int,
    aa: bool,
) -> np.ndarray:
    """Linear-unmixing decomposition for alpha matting."""
    n = len(images)
    h, w = verify_image_sizes(images)

    stacked = np.stack(images, axis=0).astype(np.float64)
    bg_arr  = np.array(bg_colors, dtype=np.float64).reshape(n, 1, 1, 3)

    out   = np.zeros((h, w, 4), dtype=np.uint8)
    alpha = np.zeros((h, w), dtype=np.float32)

    for y in range(h):
        for x in range(w):
            pxs = stacked[:, y, x, :]  # (n, 3)
            d2b = np.array([color_distance(tuple(pxs[i]), bg_colors[i]) for i in range(n)])

            # --- transparent ---
            if np.all(d2b <= tol):
                continue

            # --- opaque ---
            max_diff = max(color_distance(tuple(pxs[0]), tuple(pxs[i])) for i in range(1, n)) if n > 1 else 0.0
            if max_diff <= sim_tol:
                alpha[y, x] = 255.0
                out[y, x, :3] = pxs[0].astype(np.uint8)
                continue

            # --- edge: iterative solve ---
            s_est = pxs[int(np.argmax(d2b))].copy()
            a_est = 0.5

            for _ in range(3):
                # Solve for S given a
                w_sum = np.zeros(3)
                w_tot = 0.0
                for i in range(n):
                    if a_est <= 0.01:
                        continue
                    bg_part = (1 - a_est) * bg_arr[i, 0, 0, :]
                    s_part = (pxs[i] - bg_part) / a_est
                    wt = min(1.0, d2b[i] / tol) if tol > 0 else 1.0
                    w_sum += s_part * wt
                    w_tot += wt
                if a_est > 0.01 and w_tot > 0:
                    s_est = np.clip(w_sum / w_tot, 0, 255)

                # Solve for a given S
                alphas: list[float] = []
                for i in range(n):
                    diff = s_est - bg_arr[i, 0, 0, :]
                    dn2 = float(np.dot(diff, diff))
                    if dn2 > 1.0:
                        ai = float(np.dot(pxs[i] - bg_arr[i, 0, 0, :], diff) / dn2)
                        alphas.append(np.clip(ai, 0.0, 1.0))
                if alphas:
                    a_est = float(np.mean(alphas))

            a_final = np.clip(a_est, 0.0, 1.0)
            out[y, x, :3] = np.clip(s_est, 0, 255).astype(np.uint8)
            alpha[y, x] = a_final * 255.0 if aa else (255.0 if a_final > 0.5 else 0.0)

    out[:, :, 3] = alpha.astype(np.uint8)
    return out


# =============================================================================
# Public API — dispatcher
# =============================================================================

_METHODS = {
    "simple":        _method_simple,
    "variance":      _method_variance,
    "decomposition": _method_decomposition,
}


def remove_background(
    images: list[np.ndarray],
    bg_colors: list[Tuple[int, int, int]],
    tolerance: float,
    similarity_tolerance: Optional[float] = None,
    reference_image_idx: int = 0,
    antialiasing: bool = True,
    method: str = "simple",
) -> np.ndarray:
    """
    Remove solid backgrounds from 2-5 same-size images.

    Parameters
    ----------
    images : list[np.ndarray]
        RGB arrays, all identical (H, W), length in [2, 5].
    bg_colors : list[Tuple[int,int,int]]
        One (R,G,B) per image, in matching order.
    tolerance : float
        Max RGB distance to call a pixel "background".
    similarity_tolerance : float
        Max inter-image distance to call a pixel "subject". Defaults to *tolerance*.
    reference_image_idx : int
        Which image supplies the default RGB (for methods that don't re-colour).
    antialiasing : bool
        Smooth the alpha channel at the silhouette edge.
    method : str
        One of ``"simple"``, ``"variance"``, ``"decomposition"``.

    Returns
    -------
    np.ndarray  uint8 RGBA of shape (H, W, 4).
    """
    if similarity_tolerance is None:
        similarity_tolerance = tolerance

    fn = _METHODS.get(method)
    if fn is None:
        raise ValueError(f"Unknown method '{method}'; choose from {list(_METHODS)}")

    n = len(images)
    if not (MIN_IMAGES <= n <= MAX_IMAGES):
        raise ValueError(f"Expected {MIN_IMAGES}-{MAX_IMAGES} images, got {n}")
    if len(bg_colors) != n:
        raise ValueError(f"Need {n} bg colours, got {len(bg_colors)}")

    return fn(images, bg_colors, tolerance, similarity_tolerance, reference_image_idx, antialiasing)


# =============================================================================
# Image-to-colour matching
# =============================================================================

def _sample_border(img: np.ndarray, step: int = 1) -> np.ndarray:
    """Sample border pixels (inset from corners), return (N, 3) RGB."""
    hh, ww = img.shape[:2]
    inset = max(1, min(10, ww // 20, hh // 20))
    parts = [
        img[inset, inset : ww - inset],
        img[hh - 1 - inset, inset : ww - inset],
        img[inset : hh - inset, inset],
        img[inset : hh - inset, ww - 1 - inset],
    ]
    return np.vstack([p.reshape(-1, 3)[::step] for p in parts])


def dominant_border_color(img: np.ndarray) -> Tuple[int, int, int]:
    """Median colour of border pixels — robust estimate of the background."""
    border = _sample_border(img, step=2)
    if border.size == 0:
        return (128, 128, 128)
    m = np.median(border, axis=0)
    return tuple(int(round(c)) for c in m)  # type: ignore[return-value]


def _greedy_assign(
    unmatched: list[Tuple[int, np.ndarray, Path]],
    colors: list[Tuple[int, int, int]],
    tol: float,
    debug: bool,
) -> list[Tuple[Tuple[int, int, int], int, Path]]:
    """Greedy 1-to-1 assignment of images to colours by border-distance."""
    doms = {idx: dominant_border_color(img) for idx, img, _ in unmatched}

    pairs: list[Tuple[float, Tuple[int, int, int], int, Path]] = []
    for idx, _, path in unmatched:
        for c in colors:
            pairs.append((color_distance(doms[idx], c), c, idx, path))
    pairs.sort(key=lambda t: t[0])

    used_c: set = set()
    used_i: set = set()
    result: list[Tuple[Tuple[int, int, int], int, Path]] = []

    for dist, c, idx, path in pairs:
        if c in used_c or idx in used_i:
            continue
        used_c.add(c)
        used_i.add(idx)
        result.append((c, idx, path))
        if debug:
            print(f"    [debug] {path.name} border {doms[idx]} -> {c} (dist {dist:.1f})")
    return result


def match_images_to_colors(
    paths: list[Path],
    images: list[np.ndarray],
    candidate_colors: list[Tuple[int, int, int]],
    tol: float,
    debug: bool = False,
) -> list[Tuple[np.ndarray, Tuple[int, int, int]]]:
    """
    Order images so result[i] has background *candidate_colors[i]*.

    Strategy:
      1. Keyword in filename (white/black/red/green/blue).
      2. Dominant border colour, greedy nearest.
      3. Fallback: leftover in order (rare).
    """
    assigned: dict[Tuple[int, int, int], Tuple[int, Path]] = {}
    used: set[int] = set()

    # 1) filename
    for bg in candidate_colors:
        kws = _COLOR_KEYWORDS.get(bg, [])
        for i, p in enumerate(paths):
            if i in used:
                continue
            if any(kw in p.name.lower() for kw in kws):
                assigned[bg] = (i, p)
                used.add(i)
                print(f"  Matched {p.name} -> {bg} (filename)")
                break

    # 2) border colour
    left_paths = [(i, images[i], paths[i]) for i in range(len(images)) if i not in used]
    left_colors = [c for c in candidate_colors if c not in assigned]
    if left_paths and left_colors:
        for c, idx, p in _greedy_assign(left_paths, left_colors, tol, debug):
            assigned[c] = (idx, p)
            used.add(idx)
            print(f"  Matched {p.name} -> {c} (border colour)")

    # 3) leftovers
    still = [(i, images[i], paths[i]) for i in range(len(images)) if i not in used]
    still_c = [c for c in candidate_colors if c not in assigned]
    for j, (idx, _, p) in enumerate(still):
        if j < len(still_c):
            assigned[still_c[j]] = (idx, p)
            print(f"  Matched {p.name} -> {still_c[j]} (default order)")

    result = [(images[assigned[bg][0]], bg) for bg in candidate_colors if bg in assigned]
    if len(result) != len(images):
        raise ValueError(
            f"Could not match all {len(images)} images to colours (got {len(result)})."
        )
    return result


# =============================================================================
# CLI
# =============================================================================

def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Remove solid backgrounds from AI-generated images via pixel comparison.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
METHODS
  simple         Binary classification (fast, crisp edges).
  variance       Statistical variance model (smoother silhouettes).
  decomposition  Linear unmixing — best for semi-transparent edges,
                 drop shadows, glass, and smoke.

EXAMPLES
  # Process a folder (2-5 images)
  python remove_background.py samples/cat/

  # Explicit files
  python remove_background.py --images white.png black.png -o result.png

  # Decomposition method with custom tolerance
  python remove_background.py folder/ --method decomposition --tolerance 30
        """,
    )

    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument(
        "folder", nargs="?", type=Path,
        help=f"Folder containing {MIN_IMAGES}-{MAX_IMAGES} images.",
    )
    src.add_argument(
        "--images", nargs="*", type=Path, metavar="IMG",
        help=f"{MIN_IMAGES}-{MAX_IMAGES} image files.",
    )
    parser.add_argument("-o", "--output", type=Path, help="Output PNG path.")
    parser.add_argument(
        "--colors", nargs="*", type=str, metavar="R,G,B",
        help="Background colours as R,G,B per image.",
    )
    parser.add_argument(
        "--tolerance", type=float, default=50.0,
        help="RGB distance tolerance for background matching (default: 50).",
    )
    parser.add_argument(
        "--similarity-tolerance", type=float, default=None,
        help="Max inter-image distance for 'subject' (default = --tolerance).",
    )
    parser.add_argument(
        "--reference", type=int, default=0,
        help="0-based index of reference image for RGB output (default: 0).",
    )
    parser.add_argument(
        "--no-antialiasing", action="store_true",
        help="Hard alpha edges (0 or 255 only).",
    )
    parser.add_argument(
        "--method", default="simple", choices=list(_METHODS),
        help="Matting algorithm (default: simple).",
    )
    parser.add_argument(
        "--debug", action="store_true",
        help="Print border colours and assignment distances.",
    )
    return parser


def _parse_colors(arg: str, parser: argparse.ArgumentParser) -> Tuple[int, int, int]:
    try:
        rgb = tuple(map(int, arg.split(",")))
        if len(rgb) != 3 or not all(0 <= c <= 255 for c in rgb):
            parser.error(f"Invalid colour '{arg}': expected R,G,B with values 0-255.")
        return rgb  # type: ignore[return-value]
    except ValueError:
        parser.error(f"Invalid colour format '{arg}': expected R,G,B.")


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()

    # --- resolve paths ---
    if args.folder:
        all_paths = find_images_in_folder(args.folder)
        if len(all_paths) < MIN_IMAGES:
            parser.error(f"Folder has {len(all_paths)} image(s); need {MIN_IMAGES}-{MAX_IMAGES}.")
        if len(all_paths) > MAX_IMAGES:
            print(f"Warning: {len(all_paths)} images found; using first {MAX_IMAGES}.")
            all_paths = all_paths[:MAX_IMAGES]
        image_paths = all_paths
        out_path = args.output or args.folder / "output.png"
    else:
        image_paths = args.images or []
        if not (MIN_IMAGES <= len(image_paths) <= MAX_IMAGES):
            parser.error(f"--images needs {MIN_IMAGES}-{MAX_IMAGES} files, got {len(image_paths)}.")
        out_path = args.output or image_paths[0].parent / f"{image_paths[0].stem}_transparent.png"

    n = len(image_paths)

    # --- parse colours ---
    if args.colors:
        bg_colors = [_parse_colors(c, parser) for c in args.colors]
        if len(bg_colors) != n:
            parser.error(f"--colors has {len(bg_colors)} value(s) but there are {n} images.")
    else:
        bg_colors = DEFAULT_BG_COLORS  # matching will pick the right subset

    # --- load ---
    print(f"Loading {n} image(s)...")
    for i, p in enumerate(image_paths, 1):
        print(f"  {i}. {p.name}")

    try:
        images = load_images(image_paths)
    except Exception as exc:
        parser.error(f"Cannot load images: {exc}")

    # --- match ---
    print("\nMatching images to background colours...")
    try:
        matched = match_images_to_colors(image_paths, images, bg_colors, args.tolerance, args.debug)
        ordered = [img for img, _ in matched]
        assigned_colors = [c for _, c in matched]

        ref_img = images[args.reference]
        ref_idx = next((j for j, (img, _) in enumerate(matched) if np.array_equal(img, ref_img)), 0)
    except ValueError:
        ordered = images
        assigned_colors = DEFAULT_BG_COLORS[:n]
        ref_idx = 0
        print("  Using positional order (ensure images are correctly sorted).")

    # --- report ---
    print(f"\nBackground colours: {assigned_colors}")
    print(f"Tolerance: {args.tolerance}" + (f" | Similarity: {args.similarity_tolerance}" if args.similarity_tolerance else ""))
    print(f"Method: {args.method}")
    print(f"Reference: image_paths[{args.reference}] -> colour-order index {ref_idx}")
    print(f"Antialiasing: {'off' if args.no_antialiasing else 'on'}")
    print(f"\nProcessing ({args.method})...")

    # --- run ---
    try:
        result = remove_background(
            ordered,
            assigned_colors,
            args.tolerance,
            args.similarity_tolerance,
            ref_idx,
            antialiasing=not args.no_antialiasing,
            method=args.method,
        )
    except Exception as exc:
        parser.error(f"Processing failed: {exc}")

    Image.fromarray(result, "RGBA").save(out_path, "PNG")
    print(f"\nSaved -> {out_path}")


if __name__ == "__main__":
    main()