app.py

"""
Gradio app: object-to-object distance estimation using SAM3 + Depth Anything 3.

Deploy on Hugging Face Spaces:
  1. Create a new Space -> SDK: Gradio -> hardware: GPU recommended (SAM3 + DA3
     both run on CPU but are slow; a T4 or better is a big speedup).
  2. Upload this file as app.py, plus requirements.txt (below) and README.md.
  3. facebook/sam3 is gated: accept the license at
     https://huggingface.co/facebook/sam3, then add a `HF_TOKEN` secret to
     your Space (Settings -> Repository secrets) with a token that has access.
"""

import os
import shutil
import tempfile
import time
import numpy as np
import torch
import cv2
from PIL import Image, ExifTags
import gradio as gr

from transformers import Sam3Processor, Sam3Model
from depth_anything_3.api import DepthAnything3

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Half precision on CUDA gives a large SAM3 speed/memory win with negligible
# accuracy impact for this pipeline (mask thresholds are coarse-grained);
# CPU stays fp32 since there's no benefit there. Depth Anything 3 already
# does its own internal mixed-precision autocast (see DepthAnything3.forward
# in depth_anything_3/api.py), so it doesn't need this treatment here.
SAM3_DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
HF_TOKEN = os.environ.get("HF_TOKEN")  # set as a Space secret if sam3 is gated for you
ARUCO_DICT = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)

# --------------------------------------------------------------------------
# Lazy, cached model loading — Spaces reload this module per-worker, so we
# only want to pay the load cost once, not on every button click.
# --------------------------------------------------------------------------

_segmenter = None
_depther = None


def get_segmenter():
    global _segmenter
    if _segmenter is None:
        model = Sam3Model.from_pretrained(
            "facebook/sam3", token=HF_TOKEN, torch_dtype=SAM3_DTYPE
        ).to(DEVICE)
        processor = Sam3Processor.from_pretrained("facebook/sam3", token=HF_TOKEN)
        _segmenter = (model, processor)
    return _segmenter


def get_depther(model_id: str):
    global _depther
    if _depther is None or _depther[0] != model_id:
        model = DepthAnything3.from_pretrained(model_id).to(DEVICE)
        _depther = (model_id, model)
    return _depther[1]


# --------------------------------------------------------------------------
# Pipeline stages (same logic as the standalone script)
# --------------------------------------------------------------------------

def segment_batch(image: Image.Image, text_prompts: list, score_threshold: float = 0.5) -> list:
    """Segments every text prompt against the same primary image in a single
    batched SAM3 forward pass, instead of one full forward pass (image
    encoder included) per object as before. The image is simply repeated
    across the batch dimension so `Sam3Processor`/`Sam3Model` treat it as
    `len(text_prompts)` independent (image, text) pairs, one call instead of
    N sequential Python-level calls.

    Returns a list of (mask_or_None, error_message_or_None) tuples, one per
    entry in `text_prompts`, in the same order.
    """
    model, processor = get_segmenter()
    images = [image] * len(text_prompts)
    # dtype must match the (possibly fp16) model weights, or the forward
    # pass will fail with a dtype-mismatch error on `pixel_values`.
    inputs = processor(images=images, text=text_prompts, return_tensors="pt").to(DEVICE, dtype=SAM3_DTYPE)
    with torch.inference_mode():
        outputs = model(**inputs)

    results = processor.post_process_instance_segmentation(
        outputs,
        threshold=score_threshold,
        mask_threshold=0.5,
        target_sizes=inputs.get("original_sizes").tolist(),
    )

    per_prompt = []
    for text_prompt, result in zip(text_prompts, results):
        masks = result["masks"]
        scores = result["scores"]

        if len(masks) == 0:
            per_prompt.append((None,
                f"No object found matching '{text_prompt}' above the confidence "
                f"threshold ({score_threshold}). Try a more specific or different phrase."))
            continue

        best_idx = int(torch.argmax(scores))
        mask = masks[best_idx]
        if hasattr(mask, "cpu"):
            mask = mask.cpu().numpy()
        per_prompt.append((np.asarray(mask).astype(bool), None))
    return per_prompt


def erode_mask(mask: np.ndarray, pixels: int = 3) -> np.ndarray:
    if pixels <= 0:
        return mask
    # cv2.erode is substantially faster than scipy.ndimage.binary_erosion for
    # simple binary structuring-element erosion on 2D masks. A 3x3 cross
    # kernel matches scipy's default 4-connected structuring element.
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    eroded = cv2.erode(mask.astype(np.uint8), kernel, iterations=pixels).astype(bool)
    return eroded if eroded.sum() > 20 else mask


def robust_object_depth(depth, conf, mask, conf_percentile=40.0, mad_k=3.0):
    d, c = depth[mask], conf[mask]
    if d.size == 0:
        raise gr.Error("Mask is empty after erosion — object may be too small in this image.")

    conf_cut = np.percentile(c, conf_percentile)
    d_kept = d[c >= conf_cut]
    if d_kept.size < 5:
        d_kept = d

    med = np.median(d_kept)
    mad = np.median(np.abs(d_kept - med)) + 1e-8
    inliers = d_kept[np.abs(d_kept - med) <= mad_k * 1.4826 * mad]
    if inliers.size == 0:
        inliers = d_kept

    return float(np.median(inliers)), float(np.std(inliers)), int(inliers.size)


def backproject_to_3d(mask, robust_depth, intrinsics):
    ys, xs = np.nonzero(mask)
    cy, cx = np.median(ys), np.median(xs)
    fx, fy = intrinsics[0, 0], intrinsics[1, 1]
    px, py = intrinsics[0, 2], intrinsics[1, 2]
    z = robust_depth
    x = (cx - px) * z / fx
    y = (cy - py) * z / fy
    return np.array([x, y, z], dtype=np.float64)


OBJECT_COLORS = [
    (242, 183, 5),    # amber
    (13, 148, 136),   # teal
    (76, 29, 149),    # violet
    (224, 102, 90),   # coral
    (59, 130, 246),   # blue
    (236, 72, 153),   # pink
]


def mask_overlay(image: Image.Image, masks: list) -> Image.Image:
    """Tints each object's mask a distinct color (cycling through
    OBJECT_COLORS if there are more objects than colors) for a quick
    visual sanity-check of what got segmented."""
    # np.array(image).astype(np.float32) already allocates a fresh array, so
    # no extra .copy() is needed before mutating it in place.
    overlay = np.array(image).astype(np.float32)
    for i, mask in enumerate(masks):
        color = np.array(OBJECT_COLORS[i % len(OBJECT_COLORS)])
        overlay[mask] = overlay[mask] * 0.4 + color * 0.6
    return Image.fromarray(overlay.astype(np.uint8))


# --------------------------------------------------------------------------
# Addition 1: image-quality gate.
# A blurry/motion-blurred view doesn't just give bad depth for itself — in
# multi-view mode it can quietly drag down DA3's joint pose/depth solve for
# every other view too. Flag it before it reaches the models.
# --------------------------------------------------------------------------

def check_blur(image: Image.Image, threshold: float = 100.0):
    """Returns (sharpness_score, is_blurry). Variance of the Laplacian —
    lower means blurrier. Threshold is scene-dependent; 100 is a reasonable
    default for well-lit photos but tune it if you get false positives."""
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    score = float(cv2.Laplacian(gray, cv2.CV_64F).var())
    return score, score < threshold


# --------------------------------------------------------------------------
# Addition 2: EXIF-derived intrinsics as an alternative to DA3's estimated
# ones. If you know the real camera, this removes a whole source of error
# that no amount of downstream robust-statistics fixes.
# --------------------------------------------------------------------------

def intrinsics_from_exif(image: Image.Image):
    """Approximates fx, fy in pixels from the 35mm-equivalent focal length
    tag, assuming a 36mm-wide full-frame-equivalent sensor. Returns a 3x3
    intrinsics matrix, or None if the tag isn't present. This is an
    approximation, not a substitute for real calibration, but it's
    typically closer than a single-image network estimate."""
    try:
        exif = image.getexif()
        tag_map = {ExifTags.TAGS.get(k, k): v for k, v in exif.items()}
        focal_35mm = tag_map.get("FocalLengthIn35mmFilm")
        if not focal_35mm:
            return None
        w, h = image.size
        fx = (float(focal_35mm) / 36.0) * w
        fy = fx  # assume square pixels
        cx, cy = w / 2.0, h / 2.0
        return np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)
    except Exception:
        return None


def default_intrinsics(image: Image.Image, assumed_focal_35mm: float = 26.0) -> np.ndarray:
    """Last-resort fallback intrinsics for when neither DA3 nor the image's
    EXIF data provide any — e.g. screenshots, re-compressed/re-saved
    photos, or images from sources that strip metadata. Assumes a ~26mm
    35mm-equivalent focal length (typical of smartphone main cameras) and a
    36mm-wide full-frame-equivalent sensor, same approximation as
    `intrinsics_from_exif`. This is a coarse guess, not a calibration —
    resulting distances should be treated as rough estimates rather than
    precise measurements, but it lets the pipeline still produce a result
    instead of hard-failing."""
    w, h = image.size
    fx = (assumed_focal_35mm / 36.0) * w
    fy = fx  # assume square pixels
    cx, cy = w / 2.0, h / 2.0
    return np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)


# --------------------------------------------------------------------------
# Addition 3: automatic scale calibration via an ArUco marker of known
# physical size, instead of requiring the user to type in a measured
# reference length by hand.
#
# Print a DICT_4X4_50 marker at a known side length (e.g. 5cm) and place it
# flat in the scene. If found, this replaces the manual scale-calibration
# inputs entirely.
# --------------------------------------------------------------------------

def detect_aruco_scale(image: Image.Image, depth: np.ndarray, intrinsics: np.ndarray,
                        marker_real_size_m: float):
    if not marker_real_size_m or marker_real_size_m <= 0:
        return None, None

    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    detector = cv2.aruco.ArucoDetector(ARUCO_DICT, cv2.aruco.DetectorParameters())
    corners, ids, _ = detector.detectMarkers(gray)
    if ids is None or len(corners) == 0:
        return None, None

    quad = corners[0][0]  # 4x2 pixel corners of the first detected marker
    h, w = depth.shape
    points_3d = []
    for (px, py) in quad:
        xi, yi = int(np.clip(px, 0, w - 1)), int(np.clip(py, 0, h - 1))
        z = float(depth[yi, xi])
        fx, fy = intrinsics[0, 0], intrinsics[1, 1]
        cx, cy = intrinsics[0, 2], intrinsics[1, 2]
        x = (px - cx) * z / fx
        y = (py - cy) * z / fy
        points_3d.append(np.array([x, y, z]))

    side_lengths = [np.linalg.norm(points_3d[i] - points_3d[(i + 1) % 4]) for i in range(4)]
    measured_size = float(np.median(side_lengths))
    if measured_size <= 0:
        return None, None

    scale_correction = marker_real_size_m / measured_size
    return scale_correction, measured_size


# --------------------------------------------------------------------------
# Main entry point called by the Gradio UI
# --------------------------------------------------------------------------

def parse_object_list(objects_text: str) -> list:
    names = [n.strip() for n in objects_text.split(",")]
    names = [n for n in names if n]
    return names


def run_pipeline(files, objects_text, depth_model_id, erosion_px,
                  use_exif_intrinsics, aruco_marker_size,
                  true_reference_length, measured_reference_length):
    if not files:
        raise gr.Error("Please upload at least one image (the primary view).")

    object_names = parse_object_list(objects_text or "")
    if len(object_names) < 2:
        raise gr.Error("Please list at least two objects, comma-separated, "
                        "e.g. 'the red chair, the wooden table, the floor lamp'.")
    if len(object_names) > len(OBJECT_COLORS):
        gr.Warning(f"{len(object_names)} objects requested; overlay colors will repeat "
                   f"after the first {len(OBJECT_COLORS)}.")

    image_paths = [f.name if hasattr(f, "name") else f for f in files]
    if len(image_paths) > 5:
        gr.Warning(f"{len(image_paths)} images provided; accuracy gains from extra views "
                   f"typically plateau well before this many.")

    images = [Image.open(p).convert("RGB") for p in image_paths]

    # Addition 1: quality gate — warn (don't silently fail) on blurry views,
    # since a bad extra view can drag down the joint depth solve for all views.
    for i, img in enumerate(images):
        score, is_blurry = check_blur(img)
        if is_blurry:
            label = "primary image" if i == 0 else f"extra view {i}"
            gr.Warning(f"{label} looks blurry (sharpness score {score:.0f}). "
                       f"This can reduce accuracy — consider retaking it.")

    primary_image = images[0]

    # Segment every object in a single batched SAM3 forward pass (one image,
    # N text prompts) instead of one full forward pass per object. A bad
    # prompt for one object shouldn't discard valid results for the others,
    # so failures are collected and reported rather than raised immediately.
    masks, valid_names, seg_warnings = [], [], []
    for name, (mask, err) in zip(object_names, segment_batch(primary_image, object_names)):
        if err is not None:
            seg_warnings.append(f"'{name}': {err}")
            continue
        masks.append(erode_mask(mask, erosion_px))
        valid_names.append(name)

    for w in seg_warnings:
        gr.Warning(f"Skipped {w}")

    if len(valid_names) < 2:
        raise gr.Error("Fewer than two objects could be segmented — see warnings above for details.")

    depther = get_depther(depth_model_id)
    # DA3 only populates conf/extrinsics/intrinsics when an export step runs
    # (see https://github.com/ByteDance-Seed/Depth-Anything-3/issues/38) --
    # without export_dir/export_format they come back as None. mini_npz is
    # a lightweight raw-array export (no mesh/GLB generation), so this adds
    # negligible overhead just to force those fields to populate.
    # NOTE: DA3's export_to_mini_npz is decorated with @async_call, which just
    # fires a background Thread and returns immediately -- there's no handle to
    # join(), so inference() can return before the file is actually written.
    # Deleting the export dir right away (e.g. via a `with
    # tempfile.TemporaryDirectory()` block) races that thread and makes it crash
    # with a FileNotFoundError when it tries to write into a directory that no
    # longer exists. The intrinsics/conf fields on `prediction` itself are
    # already populated synchronously by this point, so we don't need the file
    # on disk -- we just give the background thread a brief grace period before
    # cleanup so it doesn't blow up in the Space logs.
    tmp_export_dir = tempfile.mkdtemp(prefix="da3_export_")
    try:
        prediction = depther.inference(images, export_dir=tmp_export_dir, export_format="mini_npz")
        mini_npz_path = os.path.join(tmp_export_dir, "exports", "mini_npz", "results.npz")
        for _ in range(50):  # wait up to ~5s for the async export to land
            if os.path.isfile(mini_npz_path):
                break
            time.sleep(0.1)
        exported_files = os.listdir(tmp_export_dir) if os.path.isdir(tmp_export_dir) else []
    finally:
        shutil.rmtree(tmp_export_dir, ignore_errors=True)

    # Diagnostics: these print to the Space's container logs (visible in the
    # "Logs" tab), not to the UI. If conf/intrinsics are still None after
    # requesting an export, this tells us whether the export step ran at all.
    print(f"[diag] exported files in tmp dir: {exported_files}")
    print(f"[diag] prediction attrs: {[a for a in dir(prediction) if not a.startswith('_')]}")
    print(f"[diag] depth is None: {prediction.depth is None}; "
          f"conf is None: {prediction.conf is None}; "
          f"intrinsics is None: {prediction.intrinsics is None}")

    def _to_numpy(t):
        return t.cpu().numpy() if hasattr(t, "cpu") else np.asarray(t)

    depth = np.squeeze(_to_numpy(prediction.depth[0])).astype(np.float32)

    # conf is a robustness aid, not a hard requirement -- degrade gracefully
    # to uniform confidence (median+MAD outlier rejection still applies)
    # rather than crashing if this build of DA3 doesn't return it.
    if prediction.conf is not None:
        conf = np.squeeze(_to_numpy(prediction.conf[0])).astype(np.float32)
    else:
        gr.Warning("Depth confidence map was unavailable from this DA3 build — "
                   "falling back to uniform confidence (outlier rejection still applies).")
        conf = np.ones_like(depth)

    # DA3 internally resizes images before its forward pass (see the
    # "Processed Images" log line), so depth/conf come back at that internal
    # resolution -- not the original image's. Masks (from SAM3, via
    # segment_batch's target_sizes=original_sizes) and intrinsics are both in
    # original-image pixel space, so upsample depth/conf to match before
    # indexing with the masks below (otherwise mask/depth shapes mismatch and
    # boolean indexing raises IndexError).
    target_w, target_h = primary_image.size
    if depth.shape[:2] != (target_h, target_w):
        depth = cv2.resize(depth, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
        conf = cv2.resize(conf, (target_w, target_h), interpolation=cv2.INTER_LINEAR)

    # Addition 2: prefer EXIF-derived intrinsics over DA3's estimated ones
    # when available and requested — a known camera beats a network guess.
    # Intrinsics are required for 3D back-projection; if DA3 didn't return
    # them, EXIF is tried next, and if that's also unavailable we fall back
    # to a rough default assumption (see default_intrinsics()) rather than
    # failing the whole request.
    intrinsics = prediction.intrinsics[0] if prediction.intrinsics is not None else None
    intrinsics_source = "DA3 (estimated)"
    if use_exif_intrinsics or intrinsics is None:
        exif_intrinsics = intrinsics_from_exif(primary_image)
        if exif_intrinsics is not None:
            intrinsics = exif_intrinsics
            intrinsics_source = "EXIF (35mm-equivalent focal length)"
        elif intrinsics is None:
            # Neither DA3 nor EXIF gave us intrinsics (e.g. a screenshot or a
            # re-compressed image with stripped metadata). Rather than hard-
            # failing the whole pipeline, degrade gracefully to a rough
            # default focal-length assumption so the user still gets a
            # result, just flagged as approximate.
            intrinsics = default_intrinsics(primary_image)
            intrinsics_source = "default estimate (~26mm-equivalent focal length assumption)"
            gr.Warning(
                "Camera intrinsics were unavailable from both DA3 and the image's EXIF data "
                "(e.g. no EXIF on this image). Falling back to a default focal-length "
                "assumption (~26mm-equivalent, typical smartphone camera) — treat the "
                "resulting distances as rough estimates rather than precise measurements. "
                "For better accuracy, provide a photo with intact EXIF metadata (avoid "
                "re-saving/re-compressing it, which often strips EXIF)."
            )
        else:
            gr.Warning("No usable focal-length EXIF tag found on the primary image — "
                       "falling back to DA3's estimated intrinsics.")

    # Addition 3: automatic scale calibration via ArUco marker, falling back
    # to manual true/measured length entry if no marker is found.
    scale_correction = 1.0
    scale_source = "none (raw metric depth)"
    aruco_scale, aruco_measured = detect_aruco_scale(primary_image, depth, intrinsics, aruco_marker_size)
    if aruco_scale is not None:
        scale_correction = aruco_scale
        scale_source = f"ArUco marker (measured {aruco_measured:.4f} m, expected {aruco_marker_size:.4f} m)"
    elif true_reference_length and measured_reference_length and measured_reference_length > 0:
        scale_correction = float(true_reference_length) / float(measured_reference_length)
        scale_source = "manual reference length"
    elif aruco_marker_size:
        gr.Warning("ArUco marker size was set but no marker was detected in the primary image — "
                   "check it's a DICT_4X4_50 marker, flat, and clearly visible.")

    # Per-object depth, uncertainty, and 3D point — independent of object count.
    points, stds, depths, pixel_counts = [], [], [], []
    for mask in masks:
        d, std, n = robust_object_depth(depth, conf, mask)
        p = backproject_to_3d(mask, d, intrinsics) * scale_correction
        points.append(p)
        stds.append(std * scale_correction)
        depths.append(d * scale_correction)
        pixel_counts.append(n)

    n_obj = len(valid_names)
    dist_matrix = np.zeros((n_obj, n_obj))
    unc_matrix = np.zeros((n_obj, n_obj))
    for i in range(n_obj):
        for j in range(n_obj):
            if i == j:
                continue
            dist_matrix[i, j] = np.linalg.norm(points[i] - points[j])
            unc_matrix[i, j] = np.sqrt(stds[i]**2 + stds[j]**2)

    overlay_img = mask_overlay(primary_image, masks)

    # Per-object table
    per_object_rows = "\n".join(
        f"| {name} | {depths[i]:.3f} m | {pixel_counts[i]} |"
        for i, name in enumerate(valid_names)
    )

    # Pairwise distance matrix table (upper triangle to avoid repeating each pair twice)
    header = "| |" + "".join(f" {n} |" for n in valid_names)
    sep = "|---|" + "---|" * n_obj
    rows = []
    for i in range(n_obj):
        cells = []
        for j in range(n_obj):
            if j <= i:
                cells.append(" — |")
            else:
                cells.append(f" {dist_matrix[i, j]:.3f} ± {unc_matrix[i, j]:.3f} m |")
        rows.append(f"| **{valid_names[i]}** |" + "".join(cells))
    matrix_table = "\n".join([header, sep] + rows)

    summary = (
        f"### Per-object depth\n"
        f"| Object | Depth | Pixels used |\n"
        f"|---|---|---|\n"
        f"{per_object_rows}\n\n"
        f"### Pairwise distances\n"
        f"{matrix_table}\n\n"
        f"Views used: {len(images)} "
        f"({'multi-view' if len(images) > 1 else 'single-view — add more views for better accuracy'})\n\n"
        f"Camera intrinsics: {intrinsics_source}\n\n"
        f"Scale calibration: {scale_source}"
        + (f" (×{scale_correction:.4f})" if scale_correction != 1.0 else "")
    )

    return overlay_img, summary


# --------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------

with gr.Blocks(title="Object Distance Estimator — SAM3 + Depth Anything 3") as demo:
    gr.Markdown(
        "# Object Distance Estimator\n"
        "Segment two or more objects with text prompts (SAM3), estimate metric depth "
        "(Depth Anything 3), and compute the real-world pairwise distances between them.\n\n"
        "**Tip:** upload the primary photo plus 1–4 extra photos of the *same scene* "
        "from different angles for meaningfully better accuracy."
    )

    with gr.Row():
        with gr.Column(scale=1):
            files = gr.File(
                label="Images (first = primary view, rest = optional extra views)",
                file_count="multiple",
                file_types=["image"],
            )
            objects_text = gr.Textbox(
                label="Objects (comma-separated, 2 or more)",
                placeholder="e.g. the red chair, the wooden table, the floor lamp",
            )

            with gr.Accordion("Advanced settings", open=False):
                depth_model_id = gr.Dropdown(
                    label="Depth model (must be a metric checkpoint for real distances)",
                    choices=[
                        "depth-anything/da3metric-large",
                        "depth-anything/DA3NESTED-GIANT-LARGE-1.1",
                    ],
                    value="depth-anything/da3metric-large",
                )
                erosion_px = gr.Slider(label="Mask erosion (pixels)", minimum=0, maximum=10, value=3, step=1)

                gr.Markdown("**Camera intrinsics**")
                use_exif_intrinsics = gr.Checkbox(
                    label="Prefer EXIF focal length over DA3's estimated intrinsics (if available)",
                    value=True,
                )

                gr.Markdown(
                    "**Scale calibration** — pick one: place a printed ArUco `DICT_4X4_50` "
                    "marker of known size in the scene (automatic), or enter a known length manually."
                )
                aruco_marker_size = gr.Number(label="ArUco marker side length (m)", value=None)
                true_reference_length = gr.Number(label="Manual: true length (m)", value=None)
                measured_reference_length = gr.Number(label="Manual: measured length from this pipeline (m)", value=None)

            run_btn = gr.Button("Estimate distances", variant="primary")

        with gr.Column(scale=1):
            overlay_out = gr.Image(label="Mask overlay (each object gets a distinct color)")
            result_out = gr.Markdown()

    run_btn.click(
        fn=run_pipeline,
        inputs=[files, objects_text, depth_model_id, erosion_px,
                use_exif_intrinsics, aruco_marker_size,
                true_reference_length, measured_reference_length],
        outputs=[overlay_out, result_out],
    )

if __name__ == "__main__":
    demo.launch()