fn_registry/python/functions/datascience/depth_to_relief_glb.py

"""
Construcción de una malla de relieve (heightmap) texturizada exportada como glTF binario (.glb).

Función del registry (grupo de capacidad `img-to-3d`, dominio `datascience`). Promovida desde la
app `img_to_3d_webapp`. A partir de un mapa de profundidad y la imagen original construye un grid
regular de vértices cuyo eje Z es la profundidad y mapea la imagen como textura mediante
coordenadas UV. El resultado es un modelo 3D navegable que conserva el aspecto de la imagen vista
en relieve, cargable con useGLTF / GLTFLoader directamente.

Impura: escribe el archivo .glb en disco.
"""

from __future__ import annotations

import numpy as np
from PIL import Image


def depth_to_relief_glb(
    image: "Image.Image",
    depth: "np.ndarray",
    out_glb_path: str,
    z_scale: float = 0.35,
    max_dim: int = 220,
    mask: "np.ndarray | None" = None,
) -> dict:
    """
    Construye una malla de relieve texturizada y la exporta como .glb.

    Parámetros:
        image: PIL.Image RGB usada como textura.
        depth: ndarray HxW float32 en [0,1] (1 = más cerca de la cámara).
        out_glb_path: ruta de salida del .glb.
        z_scale: amplitud del relieve (fracción del lado de la malla). Default 0.35.
        max_dim: lado máximo del grid tras downsample (controla nº de vértices/caras).
            Default 220 (~48k vértices, ~96k caras).
        mask: máscara opcional HxW (0..255, 255 = objeto), típicamente la "mask" devuelta por
            remove_background. Si se pasa, el fondo se aplana y las caras cuyos vértices caigan
            en el fondo se descartan: la malla contiene solo el objeto, sin el plano de fondo.

    Devuelve (dict, nunca lanza):
        Éxito: {"status": "ok", "glb_path": out_glb_path, "vertices": int, "faces": int,
                "height": H, "width": W}.
        Error: {"status": "error", "error": str} (depth con forma inválida, directorio de
                salida inexistente, fallo de exportación de trimesh).
    """
    try:
        import trimesh

        depth = np.asarray(depth, dtype=np.float32)
        if depth.ndim != 2:
            raise ValueError(f"depth debe ser un array 2D HxW, recibido ndim={depth.ndim}")

        H, W = depth.shape

        # Downsample para acotar el número de vértices (max_dim^2 ~ 48k vértices a 220).
        scale = max(H, W) / float(max_dim)
        if scale > 1.0:
            new_w, new_h = max(2, int(round(W / scale))), max(2, int(round(H / scale)))
            depth_img = Image.fromarray((np.clip(depth, 0, 1) * 255).astype(np.uint8))
            depth_img = depth_img.resize((new_w, new_h), Image.BILINEAR)
            depth = np.asarray(depth_img, dtype=np.float32) / 255.0
            H, W = depth.shape

        # Si se pasó máscara (objeto vs fondo), reescalarla al grid ya downsampleado: el fondo
        # no aporta relieve (se aplana a 0) y luego sus caras se descartan, dejando solo el objeto.
        fg = None
        if mask is not None:
            mask_img = Image.fromarray(np.asarray(mask).astype(np.uint8)).resize((W, H), Image.NEAREST)
            fg = np.asarray(mask_img) >= 128
            depth = np.where(fg, depth, 0.0).astype(np.float32)

        # Coordenadas del grid: X corrige aspect ratio, Y hacia abajo, Z = profundidad.
        aspect = W / float(H)
        xs = np.linspace(-aspect / 2.0, aspect / 2.0, W, dtype=np.float32)
        ys = np.linspace(0.5, -0.5, H, dtype=np.float32)
        gx, gy = np.meshgrid(xs, ys)
        gz = (depth * z_scale).astype(np.float32)
        vertices = np.column_stack([gx.ravel(), gy.ravel(), gz.ravel()])

        # Caras: dos triángulos por celda del grid.
        idx = np.arange(H * W, dtype=np.int64).reshape(H, W)
        v00 = idx[:-1, :-1].ravel()
        v01 = idx[:-1, 1:].ravel()
        v10 = idx[1:, :-1].ravel()
        v11 = idx[1:, 1:].ravel()
        faces = np.vstack(
            [
                np.column_stack([v00, v10, v11]),
                np.column_stack([v00, v11, v01]),
            ]
        )

        # Con máscara: conservar solo las caras cuyos tres vértices son objeto. La malla queda
        # recortada al objeto, sin el plano de fondo que deformaría el relieve.
        if fg is not None:
            keep = fg.ravel()[faces].all(axis=1)
            faces = faces[keep]

        # UV mapeando cada vértice al pixel de la imagen (V invertido para convención glTF).
        u = np.linspace(0.0, 1.0, W, dtype=np.float32)
        v = np.linspace(0.0, 1.0, H, dtype=np.float32)
        uu, vv = np.meshgrid(u, v)
        uv = np.column_stack([uu.ravel(), (1.0 - vv).ravel()])

        visual = trimesh.visual.TextureVisuals(uv=uv, image=image.convert("RGB"))
        mesh = trimesh.Trimesh(vertices=vertices, faces=faces, visual=visual, process=False)
        mesh.export(out_glb_path)

        return {
            "status": "ok",
            "glb_path": out_glb_path,
            "vertices": int(vertices.shape[0]),
            "faces": int(faces.shape[0]),
            "height": int(H),
            "width": int(W),
        }
    except Exception as e:  # noqa: BLE001
        return {"status": "error", "error": str(e)}


if __name__ == "__main__":
    # Demo runner end-to-end para `fn run depth_to_relief_glb_py_datascience <image_path> <out.glb>`.
    # Encadena estimate_image_depth (misma carpeta) para producir un .glb desde una imagen sin
    # tener que pasar el ndarray por CLI. La función en sí toma (image, depth); esto es solo glue
    # de demostración del flujo img→glb del grupo `img-to-3d`.
    import json
    import sys

    if len(sys.argv) < 3:
        print(json.dumps({"status": "error", "error": "uso: <image_path> <out_glb_path> [z_scale] [max_dim]"}))
        sys.exit(1)

    from estimate_image_depth import estimate_image_depth

    img_path = sys.argv[1]
    out_path = sys.argv[2]
    zs = float(sys.argv[3]) if len(sys.argv) > 3 else 0.35
    md = int(sys.argv[4]) if len(sys.argv) > 4 else 220

    est = estimate_image_depth(img_path)
    if est["status"] != "ok":
        print(json.dumps(est))
        sys.exit(1)

    res = depth_to_relief_glb(est["image"], est["depth"], out_path, z_scale=zs, max_dim=md)
    print(json.dumps(res))
    if res["status"] != "ok":
        sys.exit(1)