fn_registry/python/functions/datascience/remove_background.py

"""
Eliminación de fondo de una imagen con cascada de motores (rembg -> GrabCut -> umbral).

Función del registry (grupo de capacidad `img-to-3d`, dominio `datascience`). Promovida desde
la app `img_to_3d_webapp` (backend/bg_removal.py) para que cualquier artefacto pueda aislar el
objeto de primer plano sin reimplementar la cascada de segmentación ni la composición sobre fondo
neutro.

Impura: carga modelos neuronales (rembg/U2Net), usa GPU/CPU vía onnxruntime, lee disco y mantiene
una caché de sesión rembg a nivel de proceso para no recargar los pesos en cada llamada. Las deps
pesadas (rembg, opencv) se importan dentro de los helpers (lazy) para que el módulo se pueda
importar sin ellas; el motor de umbral es NumPy puro sin deps externas.
"""

from __future__ import annotations

import numpy as np
from PIL import Image

# Fondo gris neutro sobre el que se compone el objeto recortado.
NEUTRAL_BG = (127, 127, 127)
# Umbral de alfa para considerar un PNG RGBA "ya recortado" (passthrough).
_ALPHA_THRESH = 128
# Sesión rembg cacheada a nivel de proceso (estado mutable: ver .md "Gotchas").
_REMBG_SESSION = None


def _existing_alpha_mask(image):
    """Devuelve el canal alfa como máscara HxW uint8 si la imagen ya viene recortada, si no None."""
    if image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info):
        alpha = np.asarray(image.convert("RGBA"))[:, :, 3]
        if alpha.min() < _ALPHA_THRESH:
            return alpha
    return None


def _composite_over_neutral(image_rgb, mask):
    """Compone la imagen RGB sobre el fondo gris neutro usando la máscara como alfa."""
    rgb = np.asarray(image_rgb.convert("RGB"), dtype=np.float32)
    alpha = (mask.astype(np.float32) / 255.0)[:, :, None]
    bg = np.empty_like(rgb)
    bg[:] = NEUTRAL_BG
    out = rgb * alpha + bg * (1.0 - alpha)
    return Image.fromarray(out.clip(0, 255).astype(np.uint8), mode="RGB")


def _remove_with_rembg(image):
    """Segmenta con rembg (modelo U2Net). Devuelve (mask HxW uint8, engine_str)."""
    global _REMBG_SESSION
    from rembg import new_session, remove

    if _REMBG_SESSION is None:
        _REMBG_SESSION = new_session("u2net")
    cut = remove(image.convert("RGB"), session=_REMBG_SESSION)
    mask = np.asarray(cut.convert("RGBA"))[:, :, 3]
    return mask, "rembg:u2net"


def _remove_with_grabcut(image):
    """Segmenta con OpenCV GrabCut (rectángulo central). Devuelve (mask HxW uint8, engine_str)."""
    import cv2

    rgb = np.asarray(image.convert("RGB"))
    h, w = rgb.shape[:2]
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    gc_mask = np.zeros((h, w), np.uint8)
    bgd_model = np.zeros((1, 65), np.float64)
    fgd_model = np.zeros((1, 65), np.float64)
    margin_x, margin_y = int(0.08 * w), int(0.08 * h)
    rect = (margin_x, margin_y, max(1, w - 2 * margin_x), max(1, h - 2 * margin_y))
    cv2.grabCut(bgr, gc_mask, rect, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_RECT)
    fg = np.where((gc_mask == cv2.GC_FGD) | (gc_mask == cv2.GC_PR_FGD), 255, 0).astype(np.uint8)
    return fg, "opencv:grabcut"


def _remove_with_threshold(image):
    """Segmenta por distancia al color medio de los bordes (NumPy puro). Devuelve (mask, engine_str)."""
    rgb = np.asarray(image.convert("RGB"), dtype=np.float32)
    h, w = rgb.shape[:2]
    border = np.concatenate([rgb[0, :, :], rgb[-1, :, :], rgb[:, 0, :], rgb[:, -1, :]], axis=0)
    bg_color = border.mean(axis=0)
    dist = np.linalg.norm(rgb - bg_color, axis=2)
    thresh = max(30.0, float(dist.mean()))
    fg = (dist > thresh).astype(np.uint8) * 255
    return fg, "threshold:border"


def remove_background(image_path: str, engine: str = "auto") -> dict:
    """
    Elimina el fondo de una imagen y compone el objeto sobre un fondo gris neutro.

    Parámetros:
        image_path: ruta a la imagen de entrada (cualquier formato que PIL abra).
        engine: "auto" (default) prueba rembg -> GrabCut -> umbral en cascada y NUNCA falla
            (cae al umbral NumPy puro sin deps externas); también admite forzar un motor concreto:
            "rembg", "grabcut" o "threshold". Si se fuerza un motor y no está disponible/falla,
            o la máscara resulta degenerada, se devuelve status error.

    Devuelve (dict, nunca lanza):
        Éxito: {"status": "ok", "image": PIL.Image RGB del objeto compuesto sobre gris neutro,
                "mask": ndarray HxW uint8 (0..255, 255=objeto), "engine": str del motor usado
                ("rembg:u2net" | "opencv:grabcut" | "threshold:border" | "passthrough:alpha"),
                "height": int, "width": int, "fg_fraction": float (fracción de píxeles objeto,
                redondeada a 4 decimales)}.
        Error: {"status": "error", "error": str} (ruta inválida, motor desconocido, motor forzado
                no disponible/fallido, o ningún motor produjo una máscara válida).
    """
    try:
        image = Image.open(image_path)

        # Passthrough: si la imagen ya viene recortada (PNG RGBA con alfa), reutiliza su alfa.
        # Solo en modo auto; si se fuerza un motor concreto se respeta esa elección.
        if engine == "auto":
            existing = _existing_alpha_mask(image)
            if existing is not None:
                composed = _composite_over_neutral(image, existing)
                frac = float((existing >= 128).mean())
                h, w = existing.shape[:2]
                return {
                    "status": "ok",
                    "image": composed,
                    "mask": existing,
                    "engine": "passthrough:alpha",
                    "height": int(h),
                    "width": int(w),
                    "fg_fraction": round(frac, 4),
                }

        # Construir la lista de motores a probar según el engine pedido.
        if engine == "auto":
            attempts = [_remove_with_rembg, _remove_with_grabcut, _remove_with_threshold]
        elif engine == "rembg":
            attempts = [_remove_with_rembg]
        elif engine == "grabcut":
            attempts = [_remove_with_grabcut]
        elif engine == "threshold":
            attempts = [_remove_with_threshold]
        else:
            attempts = []

        if not attempts:
            return {"status": "error", "error": f"Motor desconocido: {engine!r}"}

        last_exc = None
        for attempt in attempts:
            try:
                mask, used = attempt(image)
            except Exception as e:  # noqa: BLE001
                last_exc = e
                continue

            # Rechazar máscaras degeneradas (casi todo fondo o casi todo objeto).
            frac = float((mask >= 128).mean())
            if frac < 0.01 or frac > 0.995:
                last_exc = f"mascara degenerada (fg_fraction={round(frac, 4)}) con {used}"
                continue

            composed = _composite_over_neutral(image, mask)
            h, w = mask.shape[:2]
            return {
                "status": "ok",
                "image": composed,
                "mask": mask,
                "engine": used,
                "height": int(h),
                "width": int(w),
                "fg_fraction": round(frac, 4),
            }

        return {
            "status": "error",
            "error": f"No se pudo eliminar el fondo con engine={engine!r}: {last_exc}",
        }
    except Exception as e:  # noqa: BLE001
        return {"status": "error", "error": str(e)}


if __name__ == "__main__":
    # Demo runner para `fn run remove_background_py_datascience <image_path> [engine] [out_dir]`.
    # Imprime un resumen JSON-serializable (el ndarray y la PIL.Image no se serializan).
    import json
    import os
    import sys

    if len(sys.argv) < 2:
        print(json.dumps({"status": "error", "error": "uso: <image_path> [engine] [out_dir]"}))
        sys.exit(1)

    path = sys.argv[1]
    eng = sys.argv[2] if len(sys.argv) > 2 else "auto"
    out_dir = sys.argv[3] if len(sys.argv) > 3 else None

    res = remove_background(path, engine=eng)
    if res["status"] == "ok":
        summary = {
            "status": "ok",
            "engine": res["engine"],
            "height": res["height"],
            "width": res["width"],
            "fg_fraction": res["fg_fraction"],
        }
        if out_dir:
            os.makedirs(out_dir, exist_ok=True)
            rgb_path = os.path.join(out_dir, "rgb.png")
            mask_path = os.path.join(out_dir, "mask.png")
            res["image"].save(rgb_path)
            Image.fromarray(res["mask"]).save(mask_path)
            summary["rgb_path"] = rgb_path
            summary["mask_path"] = mask_path
        print(json.dumps(summary))
    else:
        print(json.dumps(res))
        sys.exit(1)