app.py

import os
os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1")

import spaces  # MUST be imported before torch/cuda work
import time
import uuid
import tempfile
from pathlib import Path

import torch
import numpy as np
import cv2
import gradio as gr
from scipy.ndimage import binary_dilation, uniform_filter, uniform_filter1d
from PIL import Image
import matplotlib
import trimesh
import trimesh.visual
import utils3d
from huggingface_hub import hf_hub_download

from moge.model.v2 import MoGeModel

# ----------------------------------------------------------------------------
# Models
# ----------------------------------------------------------------------------
# DermDepth ships four checkpoints. We deliberately use two:
#
#   * metric depth + 3D  -> DermDepth_Synth_SKINL2_WoundsDB_DDI.pt
#       The paper's "DermDepth" (best model): D-Synth -> SKINL2 + WoundsDB ->
#       DDI pseudo-GT. Best metric-scale accuracy, lowest skin-tone disparity.
#
#   * surface normals    -> DermDepth_Synth_Normals.pt
#       The dedicated normal-head model. Real clinical normal ground truth is
#       noisy (SKINL2 plenoptic depth has local planar noise; WoundsDB ToF is
#       sparse and offset from RGB), so D-Synth's rendered normals are the only
#       clean normal supervision.
#
# Each checkpoint is a complete, self-contained MoGe-2 ViT-L model carrying its
# own `model_config`, and the two do NOT share a trunk -- so they are loaded as
# two instances rather than by grafting a head.
REPO_ID = "hcarrion/DermDepth"
CKPT_DEPTH = "DermDepth_Synth_SKINL2_WoundsDB_DDI.pt"
CKPT_NORMAL = "DermDepth_Synth_Normals.pt"


def load_dermdepth(filename: str) -> MoGeModel:
    path = hf_hub_download(repo_id=REPO_ID, filename=filename, repo_type="model")
    ckpt = torch.load(path, map_location="cpu", weights_only=True)
    model = MoGeModel(**ckpt["model_config"])
    model.load_state_dict(ckpt["model"], strict=False)
    return model.to("cuda").eval()


print(f"Loading metric-depth model ({CKPT_DEPTH}) ...")
model_depth = load_dermdepth(CKPT_DEPTH)
print(f"Loading normal model ({CKPT_NORMAL}) ...")
model_normal = load_dermdepth(CKPT_NORMAL)
print("Models loaded.")

# MoGe-2 turns `resolution_level` into a ViT token budget via
#     num_tokens = min + (level / 9) * (max - min),  num_tokens_range = [1200, 3600]
# Level 9 already saturates that range and nothing clamps above it, so a "level 30"
# would extrapolate to ~9200 tokens, far outside the range the model was built for.
# We address num_tokens directly and stay inside the model's real operating range.
RESOLUTION_TOKENS = {
    "Draft (1200 tokens)": 1200,
    "Balanced (2000 tokens)": 2000,
    "High (2800 tokens)": 2800,
    "Ultra (3600 tokens - max)": 3600,
}
DEFAULT_RESOLUTION = "Ultra (3600 tokens - max)"
DEPTH_CMAP = "Spectral"
# Cap rendered mesh complexity: the browser 3D viewer stalls on a full-resolution
# (~1 vertex/pixel) mesh. Measurement is unaffected -- it uses the full point map.
MESH_TARGET_VERTS = 220_000


# ----------------------------------------------------------------------------
# Visualization
# ----------------------------------------------------------------------------
def colorize_depth(depth: np.ndarray, mask=None, cmap: str = DEPTH_CMAP):
    """Colorize depth; also report the disparity range used for the mapping.

    The colormap is applied to *normalized disparity* (1/depth) -- which is why the
    colorbar's depth ticks are deliberately non-uniformly spaced.
    """
    if mask is None:
        depth = np.where(depth > 0, depth, np.nan)
    else:
        depth = np.where((depth > 0) & mask, depth, np.nan)
    disp = 1 / depth
    min_disp, max_disp = np.nanquantile(disp, 0.001), np.nanquantile(disp, 0.99)
    norm = (disp - min_disp) / (max_disp - min_disp)
    colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - norm)[..., :3], nan=0.0)
    colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8))
    return colored, float(min_disp), float(max_disp)


def render_depth_with_colorbar(colored: np.ndarray, min_disp: float, max_disp: float) -> np.ndarray:
    """Attach a metric colorbar (cm) to the colorized depth map.

    Colour is cmap(1 - t) for normalized disparity t, so the bar runs `Spectral_r`
    over t and each tick is labelled with its true depth 1/(min_disp + t*(max_disp-min_disp)).
    """
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_agg import FigureCanvasAgg
    from matplotlib.cm import ScalarMappable
    from matplotlib.colors import Normalize

    h, w = colored.shape[:2]
    if not np.isfinite([min_disp, max_disp]).all() or max_disp <= min_disp:
        return colored  # degenerate: show the raw map rather than a bogus scale

    dpi = 100
    fig = Figure(figsize=(w / dpi * 1.22, h / dpi), dpi=dpi, facecolor="white")
    FigureCanvasAgg(fig)

    ax = fig.add_axes([0.0, 0.0, 0.80, 1.0])
    ax.imshow(colored)
    ax.axis("off")

    cax = fig.add_axes([0.83, 0.06, 0.035, 0.88])
    sm = ScalarMappable(cmap=f"{DEPTH_CMAP}_r", norm=Normalize(vmin=0.0, vmax=1.0))
    cb = fig.colorbar(sm, cax=cax)

    ticks = np.linspace(0.0, 1.0, 6)
    depths_cm = [100.0 / (min_disp + t * (max_disp - min_disp)) for t in ticks]
    fmt = "{:.2f}" if max(depths_cm) < 10 else "{:.1f}"
    cb.set_ticks(ticks)
    cb.set_ticklabels([fmt.format(d) for d in depths_cm])
    cb.set_label("Metric depth (cm)", fontsize=11)
    cb.ax.tick_params(labelsize=9)
    # Do NOT invert: with t=0 at the bottom, the bar already reads
    #   top    = t=1 = Spectral_r(1) = warm = smallest depth = nearest
    #   bottom = t=0 = Spectral_r(0) = cool = largest depth  = farthest
    # which matches the captions below.
    cb.ax.text(0.5, 1.015, "near", transform=cb.ax.transAxes, ha="center", va="bottom", fontsize=8)
    cb.ax.text(0.5, -0.015, "far", transform=cb.ax.transAxes, ha="center", va="top", fontsize=8)

    fig.canvas.draw()
    return np.ascontiguousarray(np.asarray(fig.canvas.buffer_rgba())[..., :3])


def colorize_normal(normal: np.ndarray, mask=None) -> np.ndarray:
    if mask is not None:
        normal = np.where(mask[..., None], normal, 0)
    normal = normal * [0.5, -0.5, -0.5] + 0.5
    return (normal.clip(0, 1) * 255).astype(np.uint8)


# ----------------------------------------------------------------------------
# Inference
# ----------------------------------------------------------------------------
@spaces.GPU(duration=60)
def predict(
    image: np.ndarray,
    resolution_level: str = DEFAULT_RESOLUTION,
    apply_mask: bool = True,
    remove_edges: bool = True,
    max_size: int = 800,
):
    """Reconstruct metric 3D geometry from a single dermatological photograph."""
    if image is None:
        return None, None, None, "Please provide an input image.", None, None, None, [], None

    t0 = time.perf_counter()

    larger_size = max(image.shape[:2])
    if larger_size > max_size:
        scale = max_size / larger_size
        image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_AREA)

    height, width = image.shape[:2]
    num_tokens = RESOLUTION_TOKENS.get(resolution_level, 3600)

    image_tensor = torch.tensor(image, dtype=torch.float32, device="cuda").permute(2, 0, 1) / 255

    # --- metric depth + 3D, from the best model -----------------------------
    out_d = model_depth.infer(image_tensor, num_tokens=num_tokens, apply_mask=apply_mask, use_fp16=True)
    out_d = {k: v.cpu().numpy() for k, v in out_d.items()}
    points, depth, mask = out_d["points"], out_d["depth"], out_d["mask"]
    normal_geom = out_d.get("normal", None)  # same trunk as `points` -> mesh shading

    # --- surface normals, from the dedicated normal model -------------------
    out_n = model_normal.infer(image_tensor, num_tokens=num_tokens, apply_mask=apply_mask, use_fp16=True)
    out_n = {k: v.cpu().numpy() for k, v in out_n.items()}
    normal_display = out_n.get("normal", None)
    mask_n = out_n.get("mask", mask)

    mask_cleaned = mask & ~utils3d.np.depth_map_edge(depth, rtol=0.04) if remove_edges else mask

    depth_colored, min_disp, max_disp = colorize_depth(depth, mask=mask_cleaned)
    depth_vis = render_depth_with_colorbar(depth_colored, min_disp, max_disp)
    normal_vis = (
        colorize_normal(normal_display, mask=mask_n) if normal_display is not None else np.zeros_like(image)
    )

    # --- mesh (geometry + shading both from the metric model) ---------------
    # Mesh at a coarser stride than we measure at. A full-resolution point map yields
    # ~1 vertex/pixel (640k verts -> a 36 MB GLB) which locks up the browser's 3D
    # viewer. Measurement still uses the FULL-resolution `points`; only the rendered
    # geometry is decimated, and the texture stays full-res.
    stride = int(max(1, np.ceil(np.sqrt((height * width) / MESH_TARGET_VERTS))))
    pts_m = points[::stride, ::stride]
    img_m = image[::stride, ::stride]
    msk_m = mask_cleaned[::stride, ::stride]
    nrm_m = normal_geom[::stride, ::stride] if normal_geom is not None else None
    hm, wm = pts_m.shape[:2]

    if nrm_m is None:
        faces, vertices, vertex_colors, vertex_uvs = utils3d.np.build_mesh_from_map(
            pts_m, img_m.astype(np.float32) / 255, utils3d.np.uv_map(hm, wm),
            mask=msk_m, tri=True,
        )
        vertex_normals = None
    else:
        faces, vertices, vertex_colors, vertex_uvs, vertex_normals = utils3d.np.build_mesh_from_map(
            pts_m, img_m.astype(np.float32) / 255, utils3d.np.uv_map(hm, wm),
            nrm_m, mask=msk_m, tri=True,
        )

    vertices = vertices * np.array([1, -1, -1], dtype=np.float32)
    vertex_uvs = vertex_uvs * np.array([1, -1], dtype=np.float32) + np.array([0, 1], dtype=np.float32)
    if vertex_normals is not None:
        vertex_normals = vertex_normals * np.array([1, -1, -1], dtype=np.float32)

    tmpdir = Path(tempfile.gettempdir(), "dermdepth")
    tmpdir.mkdir(exist_ok=True)
    mesh_path = str(tmpdir / f"mesh_{uuid.uuid4().hex}.glb")
    trimesh.Trimesh(
        vertices=vertices, faces=faces,
        visual=trimesh.visual.texture.TextureVisuals(
            uv=vertex_uvs,
            material=trimesh.visual.material.PBRMaterial(
                baseColorTexture=Image.fromarray(image), metallicFactor=0.5, roughnessFactor=1.0,
            ),
        ),
        vertex_normals=vertex_normals, process=False,
    ).export(mesh_path)

    fov_x, fov_y = np.rad2deg(utils3d.np.intrinsics_to_fov(out_d["intrinsics"]))
    elapsed = time.perf_counter() - t0

    # DermDepth predicts *metric* depth, so these numbers carry real units.
    valid = np.isfinite(depth) & (depth > 0) & mask_cleaned
    if valid.any():
        d = depth[valid]
        d_min, d_max, d_med = float(d.min()), float(d.max()), float(np.median(d))
        scale_txt = (
            f"| **Working distance** (median depth) | **{d_med * 100:.1f} cm** |\n"
            f"| Depth range across the surface | {d_min * 100:.1f} – {d_max * 100:.1f} cm |\n"
            f"| Depth spread (1st–99th pct) | {(np.quantile(d, 0.99) - np.quantile(d, 0.01)) * 1000:.1f} mm |\n"
        )
    else:
        scale_txt = "| Metric depth | no valid depth values |\n"

    info_text = (
        "### Metric readout\n| | |\n|---|---|\n"
        f"{scale_txt}"
        f"| Field of view | {fov_x:.1f}° × {fov_y:.1f}° |\n"
        f"| Inference resolution | {num_tokens} ViT tokens |\n"
        f"| Input size | {width} × {height} px |\n"
        f"| Mesh | {(hm*wm)//1000}k verts (stride {stride}) · measured at full res |\n"
        f"| Time (both models) | {elapsed:.2f} s |\n\n"
        f"<sub>Depth & 3D from `{CKPT_DEPTH}` (best model) · normals from `{CKPT_NORMAL}`.</sub>"
    )

    # points are metric (metric_scale is applied to points and depth inside infer)
    return mesh_path, depth_vis, normal_vis, info_text, points, image, image, [], image


# ----------------------------------------------------------------------------
# Geometry: metric measurement
# ----------------------------------------------------------------------------
# SIGN CONVENTION. points[...,2] is depth: distance FROM the camera, increasing
# away from it. A lesion raised toward the lens therefore has a SMALLER z than the
# surrounding skin. We define
#       elevation = reference_z - z          (positive = raised toward the camera)
# and report "raised" and "cavity" separately.
#
# NOTE: the paper's evaluation code (fig5_ddi_volume_scatter.py,
# fig4c_ddi_lesion_measurements.py) uses `heights = Z - plane_z` and sums
# max(heights, 0) as "bump volume" -- which, under this same convention, integrates
# depressions and returns ~0 for a genuinely raised lesion. This demo deliberately
# does not mirror that.


def _xyz(points, smooth=3):
    """Point map -> X/Y/Z float64 with non-finite as NaN, lightly denoised.

    infer(apply_mask=True) sets background points to +inf; left as inf the finite
    differences below produce inf/NaN area elements that silently poison sums.

    The light box filter matters: area and volume are built from first differences,
    which are biased strictly UPWARD by per-pixel depth noise (a noisy plane has more
    apparent area than a flat one). Smoothing is applied NaN-aware so the foreground
    border does not bleed in background values.
    """
    p = np.asarray(points, dtype=np.float64).copy()
    p[~np.isfinite(p)] = np.nan
    valid = np.isfinite(p).all(axis=2)
    if smooth and smooth > 1 and valid.any():
        w = uniform_filter(valid.astype(np.float64), size=smooth, mode="nearest")
        out = np.empty_like(p)
        for k in range(3):
            ch = np.where(valid, p[..., k], 0.0)
            s = uniform_filter(ch, size=smooth, mode="nearest")
            out[..., k] = np.where(w > 1e-9, s / np.maximum(w, 1e-9), np.nan)
        out[~valid] = np.nan
        p = out
    return p[..., 0], p[..., 1], p[..., 2]


def _area_elements(X, Y, Z):
    """Return (surface_elem, proj_elem).

    surface_elem = |dP/dx x dP/dy|  -> true 3D surface area per pixel (m^2)
    proj_elem    = |nz|             -> that patch's footprint projected on the XY plane

    Volume between a surface and a reference is a column integral along Z, so its
    per-pixel weight is the PROJECTED element, not the surface element. Weighting a
    height by the surface element overestimates by 1/cos(theta) per pixel (+50% on a
    hemisphere). nz is the z-component of the same cross product.
    """
    def _fd(A, axis):
        d = np.full_like(A, np.nan)
        if axis == 1:
            d[:, :-1] = A[:, 1:] - A[:, :-1]
        else:
            d[:-1, :] = A[1:, :] - A[:-1, :]
        return d

    dXdx, dYdx, dZdx = _fd(X, 1), _fd(Y, 1), _fd(Z, 1)
    dXdy, dYdy, dZdy = _fd(X, 0), _fd(Y, 0), _fd(Z, 0)
    nx = dYdx * dZdy - dZdx * dYdy
    ny = dZdx * dXdy - dXdx * dZdy
    nz = dXdx * dYdy - dYdx * dXdy
    return np.sqrt(nx ** 2 + ny ** 2 + nz ** 2), np.abs(nz)


def _fit_reference(X, Y, Z, ring):
    """Least-squares reference surface through a ring of surrounding skin.

    A PLANE is the wrong model for healthy skin on a limb: a plane fitted around a
    patch of a 4 cm-radius forearm fabricates ~1900 mm^3 of "raised" volume where the
    truth is zero. We fit a quadric, which absorbs limb curvature, and fall back to a
    plane (then a constant) when the ring is too small to support it.

    Returns (surface_fn, rms_residual_m, model_name).
    """
    rx, ry, rz = X[ring], Y[ring], Z[ring]
    n = rx.size
    x0, y0 = float(rx.mean()), float(ry.mean())
    dx, dy = rx - x0, ry - y0

    designs = [
        ("quadric", np.column_stack([dx ** 2, dx * dy, dy ** 2, dx, dy, np.ones_like(dx)]), 12),
        ("plane", np.column_stack([dx, dy, np.ones_like(dx)]), 4),
    ]
    for name, A, need in designs:
        if n < need:
            continue
        try:
            coef, *_ = np.linalg.lstsq(A, rz, rcond=None)
        except np.linalg.LinAlgError:
            continue
        if not np.isfinite(coef).all():
            continue
        if name == "quadric":
            fn = lambda Xq, Yq, c=coef: (c[0] * (Xq - x0) ** 2 + c[1] * (Xq - x0) * (Yq - y0)
                                         + c[2] * (Yq - y0) ** 2 + c[3] * (Xq - x0)
                                         + c[4] * (Yq - y0) + c[5])
        else:
            fn = lambda Xq, Yq, c=coef: c[0] * (Xq - x0) + c[1] * (Yq - y0) + c[2]
        rms = float(np.sqrt(np.mean((rz - fn(rx, ry)) ** 2)))
        return fn, rms, name

    zc = float(np.median(rz))
    return (lambda Xq, Yq, z=zc: np.full_like(Xq, z)), float(np.std(rz)), "constant"


def compute_region_measurements(points, mask, ring_iters=None):
    """3D area / raised & cavity volume / extent for a painted region."""
    X, Y, Z = _xyz(points)
    finite = np.isfinite(X) & np.isfinite(Y) & np.isfinite(Z) & (Z > 0)
    surf_elem, proj_elem = _area_elements(X, Y, Z)
    usable = finite & np.isfinite(surf_elem) & np.isfinite(proj_elem)
    region = np.asarray(mask, bool) & usable
    n = int(region.sum())
    if n < 25:
        return {"error": "Painted region is too small, or lands on background with no valid geometry."}

    # Reference ring OUTSIDE the painted border. The paper uses `dilated & ~eroded`,
    # which straddles the border and so sits half ON the lesion -- dragging the fit
    # halfway up it (-49% on a step-bordered plateau) and making the result depend on
    # image resolution. Scale the ring to the region so it is a fixed FRACTION of it.
    if ring_iters is None:
        ring_iters = int(max(3, round(0.15 * np.sqrt(n / np.pi))))
    ring = binary_dilation(mask, iterations=ring_iters) & ~np.asarray(mask, bool) & finite
    if ring.sum() < 6:
        ring = binary_dilation(mask, iterations=ring_iters + 4) & ~np.asarray(mask, bool) & finite
    if ring.sum() < 3:
        return {"error": "No healthy skin found around the painted region to use as a reference."}

    ref_fn, rms, model = _fit_reference(X, Y, Z, ring)
    ref_z = ref_fn(X[region], Y[region])
    elevation = ref_z - Z[region]   # + = toward camera
    pe, se = proj_elem[region], surf_elem[region]

    # Extrapolation diagnostic. The ring residual (rms) only says how well the fit
    # matches the skin it SAW; it says nothing about extrapolating across the region.
    # On a healthy 4 cm forearm the ring rms is a flattering 0.05 mm while the estimate
    # still invents ~190 mm^3. What actually predicts that failure is how much the
    # reference itself BOWS across the painted area, so measure that directly: the
    # reference's own departure from a plane over the region.
    try:
        A = np.column_stack([X[region], Y[region], np.ones_like(ref_z)])
        cf, *_ = np.linalg.lstsq(A, ref_z, rcond=None)
        sag = float(np.max(np.abs(ref_z - A @ cf)))
    except np.linalg.LinAlgError:
        sag = float("nan")

    raised = float(np.sum(np.maximum(elevation, 0.0) * pe))
    cavity = float(np.sum(np.maximum(-elevation, 0.0) * pe))
    area = float(np.sum(se))

    pts = np.column_stack([X[region], Y[region], Z[region]])
    pts_c = pts - pts.mean(axis=0)
    try:
        Vt = np.linalg.svd(pts_c, full_matrices=False)[2]
        extent = float(np.ptp(pts_c @ Vt[0]))
        minor = float(np.ptp(pts_c @ Vt[1]))
    except np.linalg.LinAlgError:
        extent = minor = float("nan")

    max_raise_mm = max(float(np.nanmax(elevation)), 0.0) * 1e3
    max_depth_mm = max(float(-np.nanmin(elevation)), 0.0) * 1e3
    sag_mm = sag * 1e3
    relief_mm = max(max_raise_mm, max_depth_mm)

    # Warn only when the reference's bow is BOTH non-trivial in absolute terms and
    # comparable to the relief being claimed. Sag alone over-warns: a real 5x2 mm
    # lesion on a 4 cm forearm measures to -0.2% yet carries 0.39 mm of sag.
    warn = bool(np.isfinite(sag_mm) and sag_mm > 0.35 and sag_mm > 0.5 * relief_mm)

    return {
        "n_px": n,
        "area_mm2": area * 1e6,
        "raised_mm3": raised * 1e9,
        "cavity_mm3": cavity * 1e9,
        "max_raise_mm": max_raise_mm,
        "max_depth_mm": max_depth_mm,
        "extent_mm": extent * 1e3,
        "minor_mm": minor * 1e3,
        "ref_model": model,
        "ref_rms_mm": rms * 1e3,
        "ref_sag_mm": sag_mm,
        "curvature_warning": warn,
    }


def surface_arc_length(points, x1, y1, x2, y2, smooth=9):
    """Arc length of the surface profile along the A->B image ray.

    Not a geodesic: it follows the straight line in IMAGE space, so it can exceed the
    true shortest path over the surface.
    """
    n = int(max(abs(x2 - x1), abs(y2 - y1))) + 1
    if n < 2:
        return None
    xs = np.linspace(x1, x2, n).round().astype(int)
    ys = np.linspace(y1, y2, n).round().astype(int)
    track = np.asarray(points, dtype=np.float64)[ys, xs]
    good = np.isfinite(track).all(axis=1)
    idx = np.flatnonzero(good)
    # Refuse on gaps rather than bridging a hole with a straight chord: testing only
    # the FRACTION of good samples lets a single large hole through (a 19% gap with a
    # depth step inflated one test path by +45%).
    if idx.size < 2 or good.mean() < 0.6 or (idx.size > 1 and np.diff(idx).max() > 2):
        return None
    track = track[idx]
    if len(track) > smooth > 1:
        # Smooth only the DEVIATION from the A->B chord and pin the ends: filtering the
        # track itself pulls both endpoints inward by a constant ~2px of length, which
        # made the arc collapse onto the chord for any gently curved surface.
        t = np.linspace(0.0, 1.0, len(track))[:, None]
        base = track[0] + (track[-1] - track[0]) * t
        resid = uniform_filter1d(track - base, size=smooth, axis=0, mode="nearest")
        resid[0] = 0.0
        resid[-1] = 0.0
        track = base + resid
    return float(np.linalg.norm(np.diff(track, axis=0), axis=1).sum())


# ----------------------------------------------------------------------------
# Interactive metric measurement
# ----------------------------------------------------------------------------
MEASURE_HINT = ("Click **two points** on the image to measure the estimated metric distance "
                "between them.")


def _fmt_metric(metres: float) -> str:
    mm = metres * 1000.0
    if mm < 10:
        return f"{mm:.2f} mm"
    if mm < 1000:
        return f"{mm:.1f} mm  ({mm / 10:.2f} cm)"
    return f"{mm / 1000:.3f} m"


def _draw_marker(img, x, y, letter):
    cv2.circle(img, (x, y), 9, (255, 255, 255), -1)
    cv2.circle(img, (x, y), 9, (20, 20, 20), 2)
    cv2.putText(img, letter, (x + 13, y - 9), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 4)
    cv2.putText(img, letter, (x + 13, y - 9), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (20, 20, 20), 1)


def on_measure_click(points, base_img, clicks, evt: gr.SelectData):
    """Two clicks -> straight-line distance through metric 3D space."""
    if points is None or base_img is None:
        return None, "Run a reconstruction first, then click two points.", []

    h, w = points.shape[:2]
    x = int(np.clip(evt.index[0], 0, w - 1))
    y = int(np.clip(evt.index[1], 0, h - 1))

    clicks = list(clicks or [])
    if len(clicks) >= 2:
        clicks = []  # third click starts a fresh measurement
    clicks.append((x, y))

    img = np.ascontiguousarray(base_img.copy())
    for i, (cx, cy) in enumerate(clicks):
        _draw_marker(img, cx, cy, "AB"[i])

    if len(clicks) == 1:
        return img, "**A** set — now click point **B**.", clicks

    (x1, y1), (x2, y2) = clicks
    pa, pb = points[y1, x1], points[y2, x2]
    if not (np.isfinite(pa).all() and np.isfinite(pb).all()):
        return (
            img,
            "⚠️ One of those points has no valid geometry (it's masked background). "
            "Click on the skin surface itself.",
            clicks,
        )

    cv2.line(img, (x1, y1), (x2, y2), (255, 255, 255), 4)
    cv2.line(img, (x1, y1), (x2, y2), (20, 20, 20), 1)
    for i, (cx, cy) in enumerate(clicks):
        _draw_marker(img, cx, cy, "AB"[i])

    chord = float(np.linalg.norm(pa - pb))
    ddepth = abs(float(pa[2] - pb[2]))
    arc = surface_arc_length(points, x1, y1, x2, y2)

    rows = [f"| Straight-line (chord) | **{_fmt_metric(chord)}** |"]
    if arc is not None:
        arc = max(arc, chord)  # a surface path can never be shorter than the chord
        extra = (arc / chord - 1.0) * 100.0 if chord > 0 else 0.0
        rows.append(f"| Along the A\u2192B ray (surface profile) | {_fmt_metric(arc)}  (+{extra:.1f}% over chord) |")
    else:
        rows.append("| Along the A\u2192B ray (surface profile) | n/a \u2014 the path crosses background |")
    rows += [
        f"| Depth difference (\u0394z) | {_fmt_metric(ddepth)} |",
        f"| A (x, y, z) | {pa[0]*100:.2f}, {pa[1]*100:.2f}, {pa[2]*100:.2f} cm |",
        f"| B (x, y, z) | {pb[0]*100:.2f}, {pb[1]*100:.2f}, {pb[2]*100:.2f} cm |",
    ]
    readout = (
        f"### \U0001F4CF A \u2194 B \u2248 **{_fmt_metric(chord)}** (straight line)\n"
        "| | |\n|---|---|\n" + "\n".join(rows) + "\n\n"
        "<sub>Estimated from the predicted metric 3D point map \u2014 not pixels, and **not ground truth**. "
        "The surface profile follows the straight A\u2192B line *in the image*; it is **not** a geodesic "
        "(the shortest path over the surface), which can be shorter where it routes around raised tissue. "
        "It is lightly smoothed so per-pixel depth noise does not inflate it. The chord is the more reliable "
        "of the two. Click again to start a new measurement.</sub>"
    )
    return img, readout, clicks


def on_volume_click(editor_value, points):
    """Measure a painted region: 3D area, raised/cavity volume, width."""
    if points is None:
        return "Run a reconstruction first, then paint over a region."
    if not isinstance(editor_value, dict):
        return "Paint over the region you want to measure, then press **Measure region**."

    layers = editor_value.get("layers") or []
    mask = None
    for layer in layers:
        arr = np.asarray(layer)
        if arr.ndim != 3 or arr.shape[2] < 4:
            continue
        m = arr[..., 3] > 127
        mask = m if mask is None else (mask | m)
    if mask is None or not mask.any():
        return ("Nothing painted yet — use the brush to paint over the lesion or wound, "
                "then press **Measure region**.")

    if mask.shape != points.shape[:2]:
        # Do not stretch a mask from a different image onto this point map -- that would
        # produce confident millimetres for geometry that was never reconstructed.
        return ("The painted image no longer matches the reconstruction \u2014 press "
                "**Reconstruct in 3D** again before measuring.")

    r = compute_region_measurements(points, mask)
    if "error" in r:
        return f"⚠️ {r['error']}"

    raised, cavity = r["raised_mm3"], r["cavity_mm3"]
    hi, lo = max(raised, cavity), min(raised, cavity)
    # Only commit to a headline when one side clearly dominates: on flat, noisy skin the
    # two are both noise and the label would otherwise flip at random between runs.
    if hi < 2.0 * max(lo, 1e-9) or hi < 0.5:
        headline = "### \U0001F4D0 No clear net relief in this region"
    else:
        headline = (f"### \U0001F4D0 Estimated {'cavity' if cavity > raised else 'raised'} volume "
                    f"(vs. a fitted reference) \u2248 **{hi:,.1f} mm\u00b3**")

    warn = ""
    if r["curvature_warning"]:
        warn = ("> \u26A0\ufe0f **The surrounding skin is curved, not flat here.** The reference surface bows "
                f"{r['ref_sag_mm']:.2f} mm across your region \u2014 comparable to the relief being measured \u2014 so "
                "these volumes may be mostly body curvature rather than lesion. Paint a smaller region, or one on "
                "flatter skin.\n\n")

    return (
        f"{headline}\n\n{warn}"
        "| | |\n|---|---|\n"
        f"| Raised volume (toward camera) | {raised:,.1f} mm\u00b3 |\n"
        f"| Cavity volume (below surround) | {cavity:,.1f} mm\u00b3 |\n"
        f"| Max elevation above surround | {r['max_raise_mm']:.2f} mm |\n"
        f"| Max depth below surround | {r['max_depth_mm']:.2f} mm |\n"
        f"| 3D surface area *of painted region* | {r['area_mm2']:,.1f} mm\u00b2 ({r['area_mm2']/100:,.2f} cm\u00b2) |\n"
        f"| Longest extent *of painted region* | {r['extent_mm']:.1f} mm (minor axis {r['minor_mm']:.1f} mm) |\n"
        f"| Reference surface | {r['ref_model']} \u00b7 fit residual {r['ref_rms_mm']:.3f} mm \u00b7 bows {r['ref_sag_mm']:.2f} mm |\n"
        f"| Region size | {r['n_px']:,} px |\n\n"
        "<sub>**These are estimates against a surface least-squares fitted to a ring of skin just outside your "
        "painted border \u2014 not measurements.** Heights are weighted by each pixel's projected area, so the "
        "volume is a true column integral. **Raised** and **cavity** are separate because a nodular lesion "
        "protrudes toward the lens while an ulcer recedes from it. Volume is insensitive to how generously you "
        "paint (flat skin adds \u2248 zero), but **area and longest-extent describe the region you painted, not the "
        "lesion** \u2014 so painting past the border inflates both. Treat these as comparative (same site, same "
        "distance, over time), never absolute. Not ground truth and not a clinical measurement.</sub>"
    )


def reset_measure(base_img):
    return base_img, MEASURE_HINT, []


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
EX = "examples"
# Three held-out WoundsDB scenes (the paper splits WoundsDB by case: 1-30 train, 31+ test)
# at their native 320x240 -- which is exactly the resolution the paper evaluates WoundsDB at
# ("photo.png as the input RGB image (320x240, in thermal camera frame)"), so these are
# in-distribution rather than re-cropped. Diverse anatomy: leg, hand, foot.
WOUNDSDB_EXAMPLE = f"{EX}/woundsdb_case45_leg_venous_ulcer.png"
EXAMPLES = [
    [WOUNDSDB_EXAMPLE],
    [f"{EX}/woundsdb_case33_hand_wound.png"],
    [f"{EX}/woundsdb_case42_foot_wound.png"],
    [f"{EX}/dsynth_sample000275_fitz1-2_dark_lesion_light_skin.png"],
    [f"{EX}/dsynth_sample001300_fitz3-4_small_dark_lesion.png"],
    [f"{EX}/dsynth_sample001925_fitz5-6_multiple_lesions.png"],
]

CSS = """
#col-container { max-width: 1280px; margin: 0 auto; }
.disclaimer { border-left: 3px solid #e11d48; padding-left: 12px; }
"""

with gr.Blocks(title="DermDepth") as demo:
    points_state = gr.State(None)
    base_state = gr.State(None)
    clicks_state = gr.State([])

    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🩺 DermDepth — Monocular Metric-Scale 3D for Dermatology

            Dermatology is largely a **measurement** problem: clinicians screen and monitor lesions
            and wounds by tracking size, border, elevation and texture over time. Those properties are
            inherently 3D — yet point-of-care imaging is almost always a single 2D photo.

            **DermDepth** recovers **metric-scale** 3D from *one* ordinary photograph — no depth sensor,
            no second view, no ruler in frame. A 2.1M-parameter scale-and-normal head sits on a frozen
            [MoGe-2](https://huggingface.co/Ruicheng/moge-2-vitl-normal) backbone, trained progressively
            on **D-Synth** (synthetic renders with pixel-perfect depth, normals and intrinsics) and then
            on real clinical data. On the paper's held-out benchmarks it cuts metric scale error from
            **16.1× to 1.15×** on SKINL2 and from **81× to 1.95×** on DDI, and reduces Fitzpatrick
            skin-tone scale disparity from **10.90 to 1.02**. Those are benchmark figures — accuracy on
            your own photograph, from an unfamiliar camera or distance, may be substantially worse.

            Reconstruct an image, then use **📏 Measure distance** to click two points and read the
            estimated metric distance between them.
            """
        )
        gr.Markdown(
            "⚠️ **Research demonstration only — not a medical device.** These outputs are not "
            "diagnostic and must not inform clinical decisions. Every distance, area and volume shown is a "
            "**model estimate from a single photograph**, not a measurement — treat them as comparative, "
            "never absolute. Predictions on out-of-distribution images can fail silently.",
            elem_classes="disclaimer",
        )

        with gr.Row():
            with gr.Column(scale=4):
                input_image = gr.Image(
                    type="numpy", image_mode="RGB", label="Skin image",
                    height=340, value=WOUNDSDB_EXAMPLE,
                )
                run_btn = gr.Button("Reconstruct in 3D", variant="primary", size="lg")
                with gr.Accordion("Advanced settings", open=False):
                    resolution_level = gr.Dropdown(
                        choices=list(RESOLUTION_TOKENS.keys()), value=DEFAULT_RESOLUTION,
                        label="Inference resolution",
                        info="ViT token budget. The model's usable range is 1200–3600; Ultra is its true maximum.",
                    )
                    max_size_input = gr.Number(
                        value=800, label="Max input size (px)", precision=0, minimum=256, maximum=2048,
                        info="Longest side before inference. Drives mesh density: 800 ≈ 37 MB GLB, "
                             "1024 ≈ 60 MB (finer, slower to load).",
                    )
                    apply_mask_cb = gr.Checkbox(value=True, label="Apply predicted foreground mask")
                    remove_edges_cb = gr.Checkbox(value=True, label="Remove occlusion edges from mesh")
                info_output = gr.Markdown()

            with gr.Column(scale=6):
                with gr.Tabs():
                    with gr.Tab("🧊 3D reconstruction"):
                        mesh_output = gr.Model3D(
                            label="Drag to rotate · scroll to zoom", display_mode="solid",
                            clear_color=[0.07, 0.09, 0.12, 1.0], height=620, zoom_speed=1.2,
                        )
                    with gr.Tab("🌈 Metric depth"):
                        depth_output = gr.Image(
                            type="numpy", label="Metric depth, with scale in cm",
                            format="png", interactive=False, height=620,
                        )
                    with gr.Tab("🧭 Surface normals"):
                        normal_output = gr.Image(
                            type="numpy", label="Surface normals — from the normal-head checkpoint",
                            format="png", interactive=False, height=620,
                        )
                    with gr.Tab("📏 Measure distance"):
                        measure_image = gr.Image(
                            type="numpy", label="Click two points", format="png",
                            interactive=False, height=560,
                        )
                        measure_out = gr.Markdown(MEASURE_HINT)
                        reset_btn = gr.Button("Clear measurement", size="sm")
                    with gr.Tab("📐 Measure volume"):
                        volume_editor = gr.ImageEditor(
                            type="numpy", label="Paint over the lesion / wound",
                            brush=gr.Brush(colors=["#00e5ff"], color_mode="fixed", default_size=28),
                            eraser=gr.Eraser(default_size=28),
                            layers=False, height=520, interactive=True,
                            transforms=(), sources=(),
                        )
                        volume_btn = gr.Button("Measure region", variant="primary")
                        volume_out = gr.Markdown(
                            "Paint over the lesion or wound — keeping the **border on healthy skin** — "
                            "then press **Measure region**."
                        )

        examples_ui = gr.Examples(
            examples=EXAMPLES,
            inputs=[input_image],
            cache_examples=False,
            label="Examples — first three are real clinical photos (WoundsDB, held-out cases); last three are synthetic renders (D-Synth)",
        )

        gr.Markdown(
            """
            ---
            ### How it works

            | Output | Checkpoint | Why |
            |---|---|---|
            | Metric depth + 3D mesh | `DermDepth_Synth_SKINL2_WoundsDB_DDI.pt` | The paper's best model — D-Synth → SKINL2 + WoundsDB → DDI pseudo-GT for metric scale. |
            | Surface normals | `DermDepth_Synth_Normals.pt` | Normal-head model trained on D-Synth, whose rendered normals are the only clean normal supervision (real ToF/plenoptic normals are noisy). |

            **📏 Measure distance** reports the estimated chord and along-the-surface arc between two
            points; **📐 Measure volume** estimates 3D area and raised/cavity volume for a painted region.
            Both read through the predicted metric point map, so they are estimates in millimetres
            rather than pixel counts — not ground truth.

            ### Example credits

            The **first three** are real clinical photographs from **WoundsDB** (Chronic Wounds
            Multimodal Image Database, Silesian University of Technology), used under
            [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/) at their native 320×240 — the
            resolution the paper evaluates WoundsDB at. They are held-out cases (the paper splits
            WoundsDB by case: 1–30 train, 31+ test): `case_45` (leg), `case_33` (hand), `case_42` (foot).

            > Kręcichwost, M., Czajkowska, J., Wijata, A., Juszczyk, J., Pyciński, B., Biesok, M.,
            > Rudzki, M., Majewski, J., Kostecki, J., & Pietka, E. (2021). Chronic wounds multimodal
            > image database. *Computerized Medical Imaging and Graphics*, 88, 101844.
            > [doi:10.1016/j.compmedimag.2020.101844](https://doi.org/10.1016/j.compmedimag.2020.101844)

            The **last three** are **synthetic renders** from
            [D-Synth](https://huggingface.co/datasets/hcarrion/D-Synth) (Carrión & Norouzi),
            [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/) — one per Fitzpatrick group
            (I–II, III–IV, V–VI). They are renders, not patient photographs, and imply no diagnosis.

            No DDI imagery is bundled: Stanford's Research Use Agreement prohibits redistributing any
            portion of that dataset.

            ### Links

            📄 [Paper (MICCAI 2026)](https://arxiv.org/abs/2607.13010) ·
            🤗 [Model](https://huggingface.co/hcarrion/DermDepth) ·
            📊 [D-Synth dataset](https://huggingface.co/datasets/hcarrion/D-Synth) ·
            💻 [Code](https://github.com/hectorcarrion/dermdepth)
            """
        )

    run_btn.click(
        fn=predict,
        inputs=[input_image, resolution_level, apply_mask_cb, remove_edges_cb, max_size_input],
        outputs=[mesh_output, depth_output, normal_output, info_output,
                 points_state, base_state, measure_image, clicks_state, volume_editor],
    )
    examples_ui.load_input_event.then(
        fn=predict,
        inputs=[input_image, resolution_level, apply_mask_cb, remove_edges_cb, max_size_input],
        outputs=[mesh_output, depth_output, normal_output, info_output,
                 points_state, base_state, measure_image, clicks_state, volume_editor],
    )
    volume_btn.click(fn=on_volume_click, inputs=[volume_editor, points_state], outputs=[volume_out])
    measure_image.select(
        fn=on_measure_click,
        inputs=[points_state, base_state, clicks_state],
        outputs=[measure_image, measure_out, clicks_state],
    )
    reset_btn.click(fn=reset_measure, inputs=[base_state], outputs=[measure_image, measure_out, clicks_state])

if __name__ == "__main__":
    # Gradio 6 moved theme/css from the Blocks constructor to launch().
    demo.launch(mcp_server=True, theme=gr.themes.Default(primary_hue="teal"), css=CSS)