136 lines
4.8 KiB
Python
136 lines
4.8 KiB
Python
"""
|
|
Estimación de profundidad monocular a partir de una sola imagen con Depth-Anything-V2.
|
|
|
|
Función del registry (grupo de capacidad `img-to-3d`, dominio `datascience`). Promovida desde
|
|
la app `img_to_3d_webapp` para que cualquier artefacto pueda estimar un mapa de profundidad sin
|
|
reimplementar la carga del modelo HuggingFace ni la normalización del resultado.
|
|
|
|
Impura: descarga/carga pesos de un modelo de transformers, usa GPU si está disponible y mantiene
|
|
una caché de pipelines a nivel de proceso para no recargar en cada llamada.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
# El pipeline de transformers es caro de instanciar (carga de pesos). Se cachea por
|
|
# (modelo, device) a nivel de módulo para que un servicio no recargue en cada request.
|
|
# Es estado mutable de PROCESO: documentado como impureza (ver .md "Gotchas"). Se puede
|
|
# desactivar por llamada con use_cache=False.
|
|
_PIPE_CACHE: dict = {}
|
|
|
|
|
|
def _resolve_device(device: str) -> int:
|
|
"""Resuelve el índice de device para transformers.pipeline (0=GPU0, -1=CPU)."""
|
|
import torch
|
|
|
|
if device == "cpu":
|
|
return -1
|
|
if device == "auto":
|
|
return 0 if torch.cuda.is_available() else -1
|
|
# device explícito tipo "cuda:0" o un índice
|
|
try:
|
|
return int(device)
|
|
except ValueError:
|
|
return 0 if device.startswith("cuda") else -1
|
|
|
|
|
|
def _build_pipe(model_name: str, device: str):
|
|
from transformers import pipeline
|
|
|
|
return pipeline("depth-estimation", model=model_name, device=_resolve_device(device))
|
|
|
|
|
|
def _get_pipe(model_name: str, device: str, use_cache: bool):
|
|
if not use_cache:
|
|
return _build_pipe(model_name, device)
|
|
key = (model_name, device)
|
|
pipe = _PIPE_CACHE.get(key)
|
|
if pipe is None:
|
|
pipe = _build_pipe(model_name, device)
|
|
_PIPE_CACHE[key] = pipe
|
|
return pipe
|
|
|
|
|
|
def estimate_image_depth(
|
|
image_path: str,
|
|
model_name: str = "depth-anything/Depth-Anything-V2-Small-hf",
|
|
device: str = "auto",
|
|
use_cache: bool = True,
|
|
) -> dict:
|
|
"""
|
|
Estima un mapa de profundidad monocular a partir de una única imagen.
|
|
|
|
Parámetros:
|
|
image_path: ruta a la imagen de entrada (cualquier formato que PIL abra).
|
|
model_name: id de modelo HuggingFace de estimación de profundidad.
|
|
device: "auto" (GPU si hay), "cpu", o índice/cadena cuda explícita ("cuda:0", "0").
|
|
use_cache: si True (default) reutiliza el pipeline cacheado por (modelo, device) a
|
|
nivel de proceso; si False construye uno nuevo y no toca la caché.
|
|
|
|
Devuelve (dict, nunca lanza):
|
|
Éxito: {"status": "ok", "depth": ndarray HxW float32 normalizado a [0,1]
|
|
(1 = más cerca de la cámara), "image": PIL.Image RGB original,
|
|
"height": H, "width": W, "model": model_name, "device": device}.
|
|
Error: {"status": "error", "error": str} (ruta inválida, modelo no disponible,
|
|
device inválido, fallo de inferencia).
|
|
"""
|
|
try:
|
|
image = Image.open(image_path).convert("RGB")
|
|
pipe = _get_pipe(model_name, device, use_cache)
|
|
result = pipe(image)
|
|
depth = np.asarray(result["depth"], dtype=np.float32)
|
|
|
|
# Normalizar a [0,1]. Depth-Anything devuelve disparidad relativa (mayor = más cerca).
|
|
d = depth - depth.min()
|
|
peak = d.max()
|
|
if peak > 0:
|
|
d = d / peak
|
|
|
|
H, W = d.shape
|
|
return {
|
|
"status": "ok",
|
|
"depth": d,
|
|
"image": image,
|
|
"height": int(H),
|
|
"width": int(W),
|
|
"model": model_name,
|
|
"device": device,
|
|
}
|
|
except Exception as e: # noqa: BLE001
|
|
return {"status": "error", "error": str(e)}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Demo runner para `fn run estimate_image_depth_py_datascience <image_path> [model] [device]`.
|
|
# Imprime un resumen JSON-serializable (el ndarray y la PIL.Image no se serializan).
|
|
import json
|
|
import sys
|
|
|
|
if len(sys.argv) < 2:
|
|
print(json.dumps({"status": "error", "error": "uso: <image_path> [model_name] [device]"}))
|
|
sys.exit(1)
|
|
|
|
path = sys.argv[1]
|
|
model = sys.argv[2] if len(sys.argv) > 2 else "depth-anything/Depth-Anything-V2-Small-hf"
|
|
dev = sys.argv[3] if len(sys.argv) > 3 else "auto"
|
|
|
|
res = estimate_image_depth(path, model_name=model, device=dev)
|
|
if res["status"] == "ok":
|
|
depth = res["depth"]
|
|
summary = {
|
|
"status": "ok",
|
|
"height": res["height"],
|
|
"width": res["width"],
|
|
"depth_min": float(depth.min()),
|
|
"depth_max": float(depth.max()),
|
|
"depth_mean": round(float(depth.mean()), 4),
|
|
"model": res["model"],
|
|
"device": res["device"],
|
|
}
|
|
print(json.dumps(summary))
|
|
else:
|
|
print(json.dumps(res))
|
|
sys.exit(1)
|