app.py

"""Gradio demo for Map-Det3D."""

from __future__ import annotations

import json
import os
import re
import shutil
import tempfile
import time
import uuid
import zipfile

import gradio as gr
import numpy as np
import rerun as rr
import spaces
import torch
from gradio_rerun import Rerun
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps

from mapdet3d.model.mapanything import MapAnything
from mapdet3d.model.mapdet3d import MapDet3D, MapDet3DOut
from mapdet3d.op.mapdet3d.head import RoI2Det
from mapdet3d.vis.rerun import RerunVisualizer

HERE = os.path.dirname(os.path.abspath(__file__))
EXAMPLES_ROOT = os.path.join(HERE, "examples")
OUTPUT_ROOT = os.path.join(tempfile.gettempdir(), "mapdet3d-demo")

IMAGE_SUFFIXES = (".jpg", ".jpeg", ".png")

# Frames are logged at this long side so the recording stays web sized. The
# model always runs on the full resolution image.
VIS_LONG_SIDE = 640

MAX_FRAMES = 32
KEEP_RECORDINGS = 8
MAX_UPLOAD_BYTES = 512 * 1024 * 1024
MAX_UPLOAD_MEMBERS = 4096

# TF32
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

torch.set_float32_matmul_precision("highest")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def load_model() -> MapDet3D:
    """Load Map-Det3D without fetching the weights it immediately discards.

    ``MapDet3D.__init__`` builds its geometry backbone with
    ``MapAnything.from_pretrained``, which pulls ~5 GB of MapAnything weights
    and another ~5 GB of DINOv2 weights through torch hub. The Map-Det3D
    checkpoint then overwrites every one of them, so on a Space, where nothing
    is cached between cold starts, we build the backbone from its config alone.
    """
    config_path = hf_hub_download("facebook/map-anything", "config.json")
    with open(config_path, encoding="utf-8") as config_file:
        config = json.load(config_file)

    config["encoder_config"]["torch_hub_pretrained"] = False

    original_from_pretrained = MapAnything.from_pretrained
    MapAnything.from_pretrained = classmethod(
        lambda cls, *args, **kwargs: cls(**config)
    )
    try:
        return MapDet3D.from_pretrained("RoyYang0714/Map-Det3D")
    finally:
        MapAnything.from_pretrained = original_from_pretrained


# NOTE: ZeroGPU wants the weights placed on cuda while importing the app, a
# real GPU only shows up inside the @spaces.GPU function.
MODEL = load_model()

# Enable tracking
MODEL.track_whole_scene = True

MODEL.eval()
MODEL.to(DEVICE)


class DemoVisualizer(RerunVisualizer):
    """Rerun visualizer that logs the scene mesh once instead of per frame."""

    def __init__(self, *args, **kwargs) -> None:
        """Init."""
        super().__init__(*args, **kwargs)
        self.mesh_logged = False

    def _log_mesh(self, mesh_paths: list[str] | None) -> None:
        """Log the mesh on the first frame only, Rerun keeps it afterwards."""
        if self.mesh_logged:
            return
        super()._log_mesh(mesh_paths)
        self.mesh_logged = True


def natural_key(name: str) -> list[int | str]:
    """Sort key that orders embedded numbers by value instead of by digit."""
    return [
        int(part) if part.isdigit() else part.lower()
        for part in re.split(r"(\d+)", name)
    ]


def frame_paths(frame_dir: str, suffixes: tuple[str, ...]) -> dict[str, str]:
    """List the frames of a directory keyed by their filename stem."""
    paths: dict[str, str] = {}
    for filename in sorted(os.listdir(frame_dir)):
        stem, suffix = os.path.splitext(filename)
        if suffix.lower() in suffixes:
            paths.setdefault(stem, os.path.join(frame_dir, filename))
    return paths


def read_scene(
    scene_dir: str,
) -> tuple[list[str], dict[str, str], dict[str, str], np.ndarray, str | None]:
    """Collect the posed frames, intrinsics and mesh of a ScanNet-like scene."""
    color_dir = os.path.join(scene_dir, "color")
    pose_dir = os.path.join(scene_dir, "pose")
    intrinsic_path = os.path.join(scene_dir, "intrinsic", "intrinsic_color.txt")

    for path in (color_dir, pose_dir, intrinsic_path):
        if not os.path.exists(path):
            raise gr.Error(f"The scene is missing {os.path.relpath(path, scene_dir)}.")

    images = frame_paths(color_dir, IMAGE_SUFFIXES)
    poses = frame_paths(pose_dir, (".txt",))

    # NOTE: A pose file shares the stem of the image it belongs to.
    stems = sorted(set(images) & set(poses), key=natural_key)

    if not stems:
        raise gr.Error(
            "No frame has both a color image and a pose file of the same name."
        )

    intrinsics = np.loadtxt(intrinsic_path).astype(np.float32)[:3, :3]

    mesh_path = os.path.join(scene_dir, "mesh.ply")
    if not os.path.exists(mesh_path):
        mesh_path = None

    return stems, images, poses, intrinsics, mesh_path


def example_scenes() -> list[str]:
    """List the scenes bundled with the demo."""
    if not os.path.isdir(EXAMPLES_ROOT):
        return []
    return sorted(
        name
        for name in os.listdir(EXAMPLES_ROOT)
        if os.path.isdir(os.path.join(EXAMPLES_ROOT, name, "color"))
    )


def unpack_scene(archive_path: str) -> str:
    """Unpack an uploaded scene and return the directory that holds it."""
    if os.path.getsize(archive_path) > MAX_UPLOAD_BYTES:
        raise gr.Error(f"The archive is larger than {MAX_UPLOAD_BYTES // 1024**2} MB.")

    if not zipfile.is_zipfile(archive_path):
        raise gr.Error("Please upload the scene as a .zip archive.")

    dest = tempfile.mkdtemp(prefix="scene-", dir=OUTPUT_ROOT)

    with zipfile.ZipFile(archive_path) as archive:
        members = archive.infolist()

        if len(members) > MAX_UPLOAD_MEMBERS:
            raise gr.Error(f"The archive holds more than {MAX_UPLOAD_MEMBERS} files.")

        if sum(member.file_size for member in members) > MAX_UPLOAD_BYTES:
            raise gr.Error(
                "The archive expands to more than " f"{MAX_UPLOAD_BYTES // 1024**2} MB."
            )

        archive.extractall(dest)

    for current, dirs, _ in os.walk(dest):
        if "color" in dirs and "pose" in dirs:
            return current

    raise gr.Error(
        "The archive needs a folder holding color/, pose/ and "
        "intrinsic/intrinsic_color.txt."
    )


def prune_recordings() -> None:
    """Drop older runs and unpacked scenes so the disk does not fill up.

    Called before a run creates its own directories, so it never removes the
    ones the current request is about to use.
    """
    if not os.path.isdir(OUTPUT_ROOT):
        return

    runs = [
        entry.path
        for entry in os.scandir(OUTPUT_ROOT)
        if entry.is_dir() and entry.name.startswith(("run-", "scene-"))
    ]

    for path in sorted(runs, key=os.path.getmtime)[:-KEEP_RECORDINGS]:
        shutil.rmtree(path, ignore_errors=True)


def load_frame(image_path: str, device: str) -> torch.Tensor:
    """Load a color frame as a [1, 3, H, W] tensor of raw intensities."""
    pil_image = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB")
    image_np = np.array(pil_image).astype(np.float32)[None]

    return torch.from_numpy(np.ascontiguousarray(image_np.transpose(0, 3, 1, 2))).to(
        device
    )


def downscale_for_vis(
    image: torch.Tensor, intrinsics: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int]]:
    """Shrink a frame and its intrinsics to keep the recording web sized."""
    height, width = image.shape[2], image.shape[3]
    scale = VIS_LONG_SIDE / max(height, width)

    if scale >= 1.0:
        return image, intrinsics, (height, width)

    vis_hw = (max(1, round(height * scale)), max(1, round(width * scale)))

    vis_image = torch.nn.functional.interpolate(image, size=vis_hw, mode="area")

    # The frustum only stays put if the intrinsics follow the resolution.
    vis_intrinsics = intrinsics.clone()
    vis_intrinsics[0] *= vis_hw[1] / width
    vis_intrinsics[1] *= vis_hw[0] / height

    return vis_image, vis_intrinsics, vis_hw


def gpu_duration(
    scene_name: str, archive: str | None, num_frames: float, *_args: object
) -> int:
    """Ask for GPU time that scales with the number of frames to process."""
    return int(min(300, 60 + 3 * int(num_frames)))


@spaces.GPU(duration=gpu_duration)
def run_mapdet3d(
    scene_name: str,
    archive: str | None,
    num_frames: float,
    score_threshold: float,
    iou_threshold: float,
    show_mesh: bool,
) -> tuple[str, str, str]:
    """Detect and track objects across a scene and return a Rerun recording."""
    os.makedirs(OUTPUT_ROOT, exist_ok=True)
    prune_recordings()

    if archive:
        scene_dir = unpack_scene(archive)
    elif scene_name:
        scene_dir = os.path.join(EXAMPLES_ROOT, scene_name)
    else:
        raise gr.Error("Pick an example scene or upload one of your own.")

    stems, images, poses, intrinsics_np, mesh_path = read_scene(scene_dir)
    stems = stems[: int(num_frames)]

    log_mesh = bool(show_mesh) and mesh_path is not None

    # Get inference device
    device = "cuda" if torch.cuda.is_available() else "cpu"

    MODEL.roi2det = RoI2Det(
        nms=True,
        score_threshold=float(score_threshold),
        iou_threshold=float(iou_threshold),
    )

    # Init visualizer
    seq_name = os.path.basename(scene_dir.rstrip(os.sep)) or "scene"
    run_dir = os.path.join(OUTPUT_ROOT, f"run-{uuid.uuid4().hex}")

    visualizer = DemoVisualizer(
        convert_to_world=True,
        log_mesh=log_mesh,
        save_to_disk=True,
        output_dir=run_dir,
        start_iter=0,
    )

    # Camera intrinsics
    intrinsics = torch.from_numpy(intrinsics_np).to(device)

    start_time = time.time()

    with torch.no_grad():
        for frame_id, stem in enumerate(stems):
            # NOTE: Frame 0 clears the streaming window and the track graph,
            # so one run never inherits the state of the previous one.
            image = load_frame(images[stem], device)

            # Load pose
            extrinsics_np = np.loadtxt(poses[stem]).astype(np.float32)
            extrinsics = torch.from_numpy(extrinsics_np).to(device)

            # Run inference
            with torch.autocast("cuda", enabled=device == "cuda", dtype=torch.bfloat16):
                predictions: MapDet3DOut = MODEL(
                    images=[image],
                    intrinsics=[intrinsics],
                    extrinsics=[extrinsics],
                    frame_ids=[frame_id],
                )

            vis_image, vis_intrinsics, vis_hw = downscale_for_vis(image, intrinsics)

            visualizer.process(
                cur_iter=frame_id,
                images=[vis_image],
                sequence_names=[seq_name],
                original_hw=[vis_hw],
                intrinsics=[vis_intrinsics],
                extrinsics=[extrinsics],
                boxes3d=predictions.boxes3d,
                scores=predictions.scores,
                track_ids=predictions.track_ids,
                mesh_paths=[mesh_path] if log_mesh else None,
            )

    elapsed = time.time() - start_time

    # The visualizer streams into the file sink, flush before serving it.
    recording = rr.get_global_data_recording()
    if recording is not None:
        recording.flush()

    rrd_path = os.path.join(run_dir, "rerun_vis", f"{seq_name}.rrd")
    if not os.path.exists(rrd_path):
        raise gr.Error("Rerun did not write a recording for this scene.")

    status = (
        f"**{seq_name}** | {len(stems)} frames in {elapsed:.1f}s "
        f"({len(stems) / max(elapsed, 1e-6):.1f} FPS on {device}) | "
        f"{len(predictions.boxes3d[0])} tracked objects"
    )

    return rrd_path, rrd_path, status


with gr.Blocks(title="Map-Det3D") as demo:
    gr.HTML("""
        <h1>Map-Det3D: Metric Feed-Forward 3D Reconstruction Prior for
        Multi-view 3D Object Detection from Streaming Inputs</h1>
        <p>
            <a href="https://github.com/cvg/Map-Det3D">🌟 GitHub Repository</a> |
            <a href="https://arxiv.org/abs/2608.12179">📄 Paper</a>
        </p>
        <div style="font-size: 16px; line-height: 1.5;">
            <p>Map-Det3D detects and tracks objects in 3D from a stream of
            posed RGB frames. Frames are fed in one by one, and every detection
            is lifted into a single metric world frame, so the boxes of a scene
            accumulate into one consistent 3D layout with persistent track
            IDs.</p>
            <p>Pick a scene, hit <strong>Run Map-Det3D</strong>, then orbit the
            3D view. The mesh is shown for context only, the model never sees
            it.</p>
            <p><strong>PLEASE NOTE:</strong> We are using ZeroGPU thanks to the
            HuggingFace community Grant. Weights are moved onto the GPU on every
            inference, which adds a little time to each run. For faster
            visualization, please consider running our demo on a local machine
            from our GitHub repository.</p>
        </div>
        """)

    scenes = example_scenes()

    with gr.Row():
        with gr.Column(scale=1):
            scene_name = gr.Dropdown(
                choices=scenes,
                value=scenes[0] if scenes else None,
                label="Example scene",
            )

            num_frames = gr.Slider(
                minimum=2,
                maximum=MAX_FRAMES,
                value=21,
                step=1,
                label="Frames",
                info="More frames cover more of the scene and take longer.",
            )

            score_threshold = gr.Slider(
                minimum=0.05,
                maximum=0.9,
                value=0.25,
                step=0.05,
                label="Score threshold",
            )

            iou_threshold = gr.Slider(
                minimum=0.1,
                maximum=0.9,
                value=0.5,
                step=0.05,
                label="NMS IoU threshold",
            )

            show_mesh = gr.Checkbox(value=True, label="Show the scene mesh")

            submit_btn = gr.Button("Run Map-Det3D", scale=1, variant="primary")

            with gr.Accordion("Bring your own scene", open=False):
                gr.Markdown(
                    "Upload a `.zip` of a folder laid out like ScanNet. It "
                    "takes precedence over the example scene above.\n"
                    "```\n"
                    "scene/\n"
                    "  color/0.jpg ...          # RGB frames\n"
                    "  pose/0.txt ...           # 4x4 camera-to-world, metric\n"
                    "  intrinsic/intrinsic_color.txt\n"
                    "  mesh.ply                 # optional, context only\n"
                    "```\n"
                    "Poses have to be metric and share the stem of the frame "
                    "they belong to."
                )
                # NOTE: No file_types filter. It rejects archives client side
                # by extension alone, and unpack_scene checks the magic bytes
                # anyway, which is the stronger test.
                archive = gr.File(label="Scene archive (.zip)", type="filepath")

        with gr.Column(scale=3):
            viewer = Rerun(
                label="3D detections and tracks",
                height=720,
                panel_states={
                    "top": "hidden",
                    "blueprint": "hidden",
                    "selection": "hidden",
                },
            )
            status = gr.Markdown()
            recording_file = gr.File(label="Rerun recording (.rrd)")

    inputs = [
        scene_name,
        archive,
        num_frames,
        score_threshold,
        iou_threshold,
        show_mesh,
    ]
    outputs = [viewer, recording_file, status]

    # NOTE: No gr.Examples here. It cannot hold the upload alongside the other
    # controls: Gradio drops a File column from the examples table, so the row
    # carried five values into six inputs and every argument after the scene
    # name shifted by one, landing the frame count where the archive goes. The
    # dropdown already selects the only example scene and the sliders already
    # default to its settings, so the table added nothing.
    submit_btn.click(fn=run_mapdet3d, inputs=inputs, outputs=outputs)

if __name__ == "__main__":
    """Demo."""
    os.makedirs(OUTPUT_ROOT, exist_ok=True)

    # NOTE: Spaces turn Gradio's server-side rendering on through
    # GRADIO_SSR_MODE, and its node server answers the browser's POSTs with
    # `405 POST method not allowed`, so nothing can be run or uploaded from the
    # page. Rendering client side keeps every request on the Python backend.
    demo.launch(allowed_paths=[OUTPUT_ROOT], ssr_mode=False)