feat(datascience): promover remove_background al registry + mask en depth_to_relief_glb (grupo img-to-3d)
Completa la promoción del flujo imagen->3D al registry (grupo de capacidad img-to-3d), extraído de la app img_to_3d_webapp. - remove_background_py_datascience (nueva): elimina el fondo con cascada rembg/U2Net -> OpenCV GrabCut -> umbral NumPy, compone el objeto sobre gris neutro y devuelve image + mask + engine. Impura, nunca lanza. Adaptada de backend/bg_removal.py con firma de ruta (image_path) y salida dict, demo CLI JSON-serializable. - depth_to_relief_glb_py_datascience (v1.1.0): añade el parámetro opcional mask para recortar la malla de relieve al objeto (descarta las caras del fondo), cerrando la cadena con remove_background. Aditivo (mask=None = comportamiento previo), fiel al original de backend/depth.py. - docs/capabilities/img-to-3d.md: incorpora remove_background como paso 0 (pre-proceso), actualiza el flujo a 3 pasos encadenados, la tabla de funciones (4), el ejemplo end-to-end con mask y las deps (rembg/opencv). - docs/capabilities/INDEX.md: conteo del grupo 3 -> 4. Las dos funciones ya presentes (estimate_image_depth, depth_to_relief_glb) y el pipeline build_relief_glb_from_image fueron promovidas en una ronda previa. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,213 @@
|
||||
"""
|
||||
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)
|
||||
Reference in New Issue
Block a user