""" 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 [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: [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)