"""ExMR: Extremity Motion Refinement for Monocular Human Motion Capture.

A ZeroGPU Gradio demo that reconstructs 3D body, hand and foot motion (SMPL-X) of the main
person in a monocular video, using the ExMR model (https://huggingface.co/papers/2610.26213).
Adapted from the official tools/demo.py of https://github.com/twehrbein/ExMR.

Note on the pipeline: the official exmr.pipeline.run_video runs SimpleVO visual odometry in a
separate PROCESS (ProcessPoolExecutor). ZeroGPU forbids spawning child processes inside
@spaces.GPU, so this app runs SimpleVO synchronously in a background THREAD instead (it is
CPU-only SIFT + COLMAP two-view geometry, so GIL contention is minimal), and keeps the static
camera path (no odometry at all) as the default, matching the authors' example video.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# SimpleVO / pycolmap must not fork; OpenCV's threads inside the ZeroGPU worker are fine
os.environ.setdefault("OMP_NUM_THREADS", "8")

import spaces  # noqa: E402  (must precede any torch / CUDA import)

import json
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE / "lib"))


def _patch_torch_dynamo_config():
    """ZeroGPU's torch 2.11.0+cu130 ships a torch._dynamo.config that passes an unknown
    `deprecated=` kwarg into torch.utils._config_module.Config, which breaks
    `import pytorch_lightning` (needed by HaMeR and SAM 3D Body). Drop unknown kwargs."""
    try:
        import torch.utils._config_module as cm

        _orig = cm.Config.__init__

        def _safe_init(self, **kwargs):
            allowed = getattr(cm.Config, "_defaults", None)
            if allowed:
                kwargs = {k: v for k, v in kwargs.items() if k in allowed}
            return _orig(self, **kwargs)

        cm.Config.__init__ = _safe_init
    except Exception:
        pass


_patch_torch_dynamo_config()

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402

import logging  # noqa: E402

log = logging.getLogger("exmr-demo")
logging.basicConfig(level=logging.INFO)

# ---------------------------------------------------------------------------
# Checkpoint download: ~12 GB of models fetched at startup into a local cache.
# Mirrors are SHA256-verified against ExMR's own tools/download_models.py pin table; the SMPL-X
# body model and the MANO hand model use the public mirrors of the registered MPS files.
# ---------------------------------------------------------------------------
CKPT_DIR = Path(os.environ.get("EXMR_CHECKPOINT_DIR", _HERE / "checkpoints"))

_DINOV3_COMMIT = "6876159a11b4df116f30f667f8c9888617df0751"
_DINOV3_TREE_SHA256 = "0389aafef4614add2090f2907f210698df48f16d3f2b845b6002c4ebb6aea0f8"


def _download_dinov3_code():
    """Pre-fill torch.hub's cache with the pinned DINOv3 commit SAM 3D Body's backbone needs.

    torch.hub.load(..., source="github") downloads the repo zip on first use; doing it here
    (from the bundled copy at startup) keeps the first GPU call from doing a 40 MB network fetch
    with weird failure modes inside the ZeroGPU worker.
    """
    import zipfile

    hub_dir = Path(torch.hub.get_dir())
    repo_dir = hub_dir / "facebookresearch_dinov3_main"
    if repo_dir.exists():
        return
    url = f"https://github.com/facebookresearch/dinov3/archive/{_DINOV3_COMMIT}.zip"
    hub_dir.mkdir(parents=True, exist_ok=True)
    tmp = hub_dir / f"_dinov3_{_DINOV3_COMMIT[:12]}.zip"
    if not tmp.exists():
        subprocess.run(
            ["curl", "-sSL", "--retry", "3", "-o", str(tmp), url], check=True
        )
    repo_dir.parent.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(tmp) as z:
        z.extractall(str(hub_dir / "_dinov3_tmp"))
    extracted = hub_dir / "_dinov3_tmp" / f"dinov3-{_DINOV3_COMMIT}"
    shutil.move(str(extracted), str(repo_dir))
    shutil.rmtree(str(hub_dir / "_dinov3_tmp"), ignore_errors=True)
    tmp.unlink(missing_ok=True)
    marker = repo_dir / "hubconf.py"
    if not marker.exists():
        raise RuntimeError("DINOv3 hub prefetch failed: hubconf.py missing")


def _fetch(name, repo, filename):
    """Download one file of an HF repo into CKPT_DIR/<name> (the exmr.assets layout)."""
    import urllib.request

    target = CKPT_DIR / name
    if target.exists():
        return target
    log.info(f"downloading {repo}/{filename} -> {target}")
    target.parent.mkdir(parents=True, exist_ok=True)
    # hf_hub_download(local_dir=CKPT_DIR) places the file at CKPT_DIR/<filename-path>;
    # rename it into the target layout (same filesystem, so a cheap os.rename).
    fetched = hf_hub_download(repo, filename, local_dir=str(CKPT_DIR))
    src = Path(fetched)
    if src != target:
        if target.exists():
            target.unlink()
        src.rename(target)
    return target


def _convert_mano_pickle(path) -> None:
    """Re-serialize hamer/data/mano/MANO_RIGHT.pkl without chumpy (in place).

    The upstream MANO v1.2 pickle stores `shapedirs` as a `chumpy.Select`, and chumpy does
    not work on Python 3.12 / numpy 2. `smplx.MANOLayer` (inside HaMeR's MANO wrapper) does a
    plain `pickle.load(mano_file, encoding='latin1')`, which therefore dies with
    `ModuleNotFoundError: No module named 'chumpy'`. The authors' own downloader
    (tools/download_models.py) converts the pickle to plain numpy/scipy and verifies the
    resulting arrays against a hash of the tested MANO v1.2 model -- we do exactly the same.
    Idempotent: a file that already loads as a plain pickle is left alone.
    """
    import hashlib
    import pickle

    import numpy as np
    import scipy.sparse

    path = Path(path)

    try:  # already converted?
        with open(path, "rb") as fh:
            pickle.load(fh, encoding="latin1")
        log.info("MANO_RIGHT.pkl already chumpy-free")
        return
    except Exception:
        pass

    class _ChumpyStub:
        """Stands in for chumpy objects (chumpy does not work on Python 3.12 / numpy 2)."""

        def __init__(self, *args, **kwargs):
            pass

        def __setstate__(self, state):
            self.state = state

    class _ManoUnpickler(pickle.Unpickler):
        ALLOWED = {
            ("numpy", "ndarray"),
            ("numpy", "dtype"),
            ("numpy.core.multiarray", "_reconstruct"),
            ("numpy._core.multiarray", "_reconstruct"),
            ("scipy.sparse.csc", "csc_matrix"),
            ("scipy.sparse._csc", "csc_matrix"),
            ("__builtin__", "set"),
            ("builtins", "set"),
            ("__builtin__", "object"),
            ("builtins", "object"),
            ("copy_reg", "_reconstructor"),
            ("copyreg", "_reconstructor"),
        }

        def find_class(self, module, name):
            if module in ("chumpy.ch", "chumpy.reordering"):
                return type(name, (_ChumpyStub,), {})
            if (module, name) not in self.ALLOWED:
                raise pickle.UnpicklingError(f"Unexpected object in the MANO model: {module}.{name}")
            if name == "csc_matrix":  # pickled under the deprecated module scipy.sparse.csc
                return scipy.sparse.csc_matrix
            return super().find_class(module, name)

    def _resolve_chumpy(obj):
        kind, state = type(obj).__name__, obj.state
        if kind == "Ch":
            return np.asarray(state["x"])
        if kind == "Select":
            a = _resolve_chumpy(state["a"])
            return a.ravel()[np.asarray(state["idxs"]).ravel()].reshape(state["preferred_shape"])
        raise RuntimeError(f"Unsupported chumpy object in the MANO model: {kind}")

    with open(path, "rb") as fh:
        model = _ManoUnpickler(fh, encoding="latin1").load()
    for key, value in model.items():
        if isinstance(value, _ChumpyStub):
            model[key] = _resolve_chumpy(value)

    digest = hashlib.sha256()
    for key in sorted(model):
        value = model[key]
        if scipy.sparse.issparse(value):
            value = value.toarray()
        value = np.ascontiguousarray(value) if isinstance(value, np.ndarray) else str(value)
        digest.update(key.encode() + b"\0")
        if isinstance(value, np.ndarray):
            digest.update(f"{value.dtype.str}{value.shape}".encode() + value.tobytes())
        else:
            digest.update(value.encode())
    actual = digest.hexdigest()
    expected = "a7a63a5c0b352fbf7e1b8a565e353e126c63066de0cc0c9f8af031dff59d98c2"
    if actual != expected:
        raise RuntimeError(
            f"MANO_RIGHT.pkl converted arrays differ from the tested MANO v1.2 model:\n"
            f"  expected {expected}\n  got      {actual}"
        )

    with open(path, "wb") as fh:
        pickle.dump(model, fh)
    log.info(f"MANO_RIGHT.pkl converted to a chumpy-free pickle (arrays sha256 {actual[:16]}...)")


def _prepare_assets() -> None:
    """Materialize the checkpoint layout expected by exmr.assets.Assets."""
    import urllib.request

    # 1) ExMR transformer weights
    _fetch("exmr/exmr.safetensors", "twehrbein/ExMR", "exmr.safetensors")
    # 2) SMPL-X neutral body model (registered model; public mirror, SHA-verified upstream)
    _fetch(
        "body_models/smplx/SMPLX_NEUTRAL.npz", "artabatest/smplx_modelv2", "smplx/SMPLX_NEUTRAL.npz"
    )
    # 3) YOLO person detectors straight from Ultralytics' release assets
    for yolo in ("yolo11m", "yolo11x"):
        target = CKPT_DIR / "yolo" / f"{yolo}.pt"
        if target.exists():
            continue
        target.parent.mkdir(parents=True, exist_ok=True)

        def _dl(name):
            url = f"https://github.com/ultralytics/assets/releases/download/v8.3.0/{name}.pt"
            with urllib.request.urlopen(
                urllib.request.Request(url, headers={"User-Agent": "exmr-demo"})
            ) as r, open(target, "wb") as f:
                shutil.copyfileobj(r, f, 1 << 20)

        _dl(yolo)
    # 4) Sapiens2 1B pose
    _fetch(
        "sapiens2/sapiens2_1b_pose.safetensors",
        "facebook/sapiens2-pose-1b",
        "sapiens2_1b_pose.safetensors",
    )
    # 5) SAM 3D Body (DINOv3-H+ backbone; weights all in model.ckpt, mhr body model separate)
    _fetch("sam-3d-body-dinov3/model.ckpt", "mingjiexie/sam-3d-body-dinov3", "model.ckpt")
    _fetch(
        "sam-3d-body-dinov3/model_config.yaml",
        "mingjiexie/sam-3d-body-dinov3",
        "model_config.yaml",
    )
    _fetch(
        "sam-3d-body-dinov3/assets/mhr_model.pt",
        "mingjiexie/sam-3d-body-dinov3",
        "assets/mhr_model.pt",
    )
    # 6) HaMeR (checkpoint + config + MANO + mean params)
    _fetch("hamer/hamer_ckpts/checkpoints/hamer.ckpt", "CometsFeiyu/hamer", "hamer_ckpts/checkpoints/hamer.ckpt")
    _fetch("hamer/hamer_ckpts/model_config.yaml", "CometsFeiyu/hamer", "hamer_ckpts/model_config.yaml")
    _fetch("hamer/hamer_ckpts/dataset_config.yaml", "CometsFeiyu/hamer", "hamer_ckpts/dataset_config.yaml")
    _fetch("hamer/data/mano_mean_params.npz", "CometsFeiyu/hamer", "data/mano_mean_params.npz")
    # The upstream MANO_RIGHT.pkl is a chumpy pickle; convert it in place before HaMeR loads it.
    _convert_mano_pickle(
        _fetch("hamer/data/mano/MANO_RIGHT.pkl", "CometsFeiyu/hamer", "data/mano/MANO_RIGHT.pkl")
    )
    # 7) the DINOv3 hub code cache
    _download_dinov3_code()
    log.info("all checkpoints ready")


_prepare_assets()

# The pipeline modules import heavy third-party code (pytorch_lightning, ultralytics, sapiens);
# import them AFTER the dynamo patch above.
from exmr.assets import Assets  # noqa: E402
from exmr.model.exmr import ExMR  # noqa: E402

ASSETS = Assets(root=CKPT_DIR)
ASSETS.check("sapiens2", "yolo11m", False)

# ---------------------------------------------------------------------------
# Model loading (module scope, eagerly on the GPU): every stage of the pipeline is loaded once
# ---------------------------------------------------------------------------
log.info("loading models...")
DEVICE = "cuda"

from exmr.body_models.smplx_model import SmplxModel  # noqa: E402
from exmr.model.exmr import ExMR as _ExMR  # noqa: F401, E402
from exmr.preprocessing.body_features import SamBodyFeatures  # noqa: E402
from exmr.preprocessing.hand_features import HamerHandFeatures  # noqa: E402
from exmr.preprocessing.keypoints.mappings import (  # noqa: E402
    COCO23,
    COCO63_LEFT_HAND,
    COCO63_RIGHT_HAND,
)
from exmr.preprocessing.keypoints import VISIBILITY_THRESHOLD  # noqa: E402
from exmr.preprocessing.keypoints.sapiens2 import Sapiens2Keypoints  # noqa: E402
from exmr.preprocessing.tracker import Tracker  # noqa: E402
from exmr.vis.render import render_videos  # noqa: E402
from exmr.geometry.camera import (  # noqa: E402
    bbx_xys_from_xyxy,
    cam_angvel_from_rotations,
    estimate_K,
)
from exmr.preprocessing import video as V  # noqa: E402

torch.set_grad_enabled(False)

# Load every GPU stage eagerly at module scope, as ZeroGPU requires: the weights get packed to
# disk once (at boot) and streamed into VRAM on each @spaces.GPU entry, instead of paying the
# multi-minute load on the first request.
log.info("loading ExMR + SMPL-X")
EXMR = ExMR.from_pretrained(ASSETS.exmr, ASSETS.smplx).to(DEVICE).eval()
SMPLX = SmplxModel(ASSETS.smplx).to(DEVICE)

log.info("loading YOLO trackers")
TRACKER = {
    "yolo11m": Tracker(ASSETS.yolo("yolo11m")),
    "yolo11x": Tracker(ASSETS.yolo("yolo11x")),
}

log.info("loading Sapiens2 (448x336)")
SAPIENS2 = {res: None for res in ("448x336", "576x432", "1024x768")}
# num_workers=0: the ZeroGPU worker process is daemonic, and DataLoader workers are
# child processes -> "daemonic processes are not allowed to have children".
SAPIENS2["448x336"] = Sapiens2Keypoints(ASSETS.sapiens2, resolution="448x336", num_workers=0)

log.info("loading SAM 3D Body")
SAM_BODY = SamBodyFeatures(ASSETS.sam_3d_body, num_workers=0)

log.info("loading HaMeR")
HAMER = HamerHandFeatures(ASSETS.hamer, num_workers=0)

MODELS_LOCK = threading.Lock()


def _get_sapiens2(resolution: str) -> Sapiens2Keypoints:
    """The 448x336 model loads at startup; other resolutions lazily on first use."""
    with MODELS_LOCK:
        if SAPIENS2.get(resolution) is None:
            SAPIENS2[resolution] = Sapiens2Keypoints(
                ASSETS.sapiens2, resolution=resolution, num_workers=0
            )
        return SAPIENS2[resolution]


# ---------------------------------------------------------------------------
# The pipeline (synchronous port of exmr.pipeline.run_video, without process pools)
# ---------------------------------------------------------------------------
def _simplevo_in_thread(video_path):
    """Run SimpleVO in a background THREAD (its numpy/pycolmap calls release the GIL)."""
    from exmr.preprocessing.visual_odometry import SimpleVO

    def work():
        try:
            result["rot"] = SimpleVO(progress=False)(video_path)
        except BaseException as e:
            result["err"] = e

    result = {}
    thread = threading.Thread(target=work, daemon=True)
    thread.start()
    return result, thread


def _run_exmr(video_path, output_dir: Path, detector, resolution, static_cam):
    """Run the full ExMR pipeline on one video; returns (results, out_paths)."""
    from exmr.preprocessing.hand_features import hand_boxes_from_keypoints

    name = Path(video_path).stem
    output_dir.mkdir(parents=True, exist_ok=True)

    frame_dir = output_dir / "frames"
    num_frames = V.extract_frames(video_path, frame_dir)
    if V.video_lwh(video_path)[0] != num_frames:
        log.warning("frame count mismatch; rebuilding constant-frame-rate video")
        cfr = output_dir / f"{name}_cfr.mp4"
        if not cfr.exists():
            V.encode_frames(frame_dir, cfr, fps=V.video_fps(video_path))
        video_path = str(cfr)
    width, height = V.frame_size(frame_dir)
    fps = V.video_fps(video_path)
    log.info(f"{video_path}: {num_frames} frames, {width}x{height}, {fps:.2f} fps")

    L = num_frames

    # SimpleVO (CPU) overlaps the GPU stages below
    vo_thread = None
    vo_result = None
    if not static_cam:
        vo_result, vo_thread = _simplevo_in_thread(video_path)

    # --- tracking (YOLO11 + ByteTrack)
    tracker = TRACKER[detector]
    bbx_xyxy = tracker.get_main_track(video_path).float()  # (L, 4)
    bbx_xys = bbx_xys_from_xyxy(bbx_xyxy, base_enlarge=1.2).float()
    K = estimate_K(width, height)[None].repeat(L, 1, 1)

    # --- 2D whole-body keypoints
    kp2d = _get_sapiens2(resolution)(frame_dir, bbx_xyxy)

    # --- SAM 3D Body image features
    body_feat = SAM_BODY(frame_dir, bbx_xyxy, K)

    # --- HaMeR hand features
    threshold = VISIBILITY_THRESHOLD["sapiens2"]
    boxes = torch.stack(
        [
            hand_boxes_from_keypoints(kp2d[:, COCO63_LEFT_HAND], threshold),
            hand_boxes_from_keypoints(kp2d[:, COCO63_RIGHT_HAND], threshold),
        ],
        dim=1,
    )
    hand_feat = HAMER(frame_dir, boxes)

    # --- camera rotations
    if static_cam:
        R_w2c = torch.eye(3).repeat(L, 1, 1)
    else:
        if vo_thread is not None:
            vo_thread.join()
        if "err" in vo_result:
            raise RuntimeError(f"visual odometry failed: {vo_result['err']}")
        R_w2c = vo_result["rot"]

    kp23 = kp2d[:, COCO23].clone()

    inputs = {
        "kp2d": kp23[None],
        "bbx_xys": bbx_xys[None],
        "K_fullimg": K[None],
        "cam_angvel": cam_angvel_from_rotations(R_w2c)[None],
        "body_feat": body_feat[None],
        "hand_feat": hand_feat[None],
    }
    with torch.no_grad():
        out = EXMR(
            **{k: v.to(DEVICE) for k, v in inputs.items()},
            static_cam=static_cam,
        )
    results = {
        "smplx_incam": {k: v[0].cpu() for k, v in out["smplx_incam"].items()},
        "smplx_global": {k: v[0].cpu() for k, v in out["smplx_global"].items()},
        "static_conf_logits": out["static_conf_logits"][0].cpu(),
        "K_fullimg": K,
        "bbx_xyxy": bbx_xyxy,
    }

    # --- render the three views
    with torch.no_grad():
        verts_incam = SMPLX(**{k: v.to(DEVICE) for k, v in results["smplx_incam"].items()})
        verts_global = SMPLX(**{k: v.to(DEVICE) for k, v in results["smplx_global"].items()})
    out_paths = [
        output_dir / f"{name}_{view}.mp4" for view in ("incam", "global", "incam_global")
    ]
    render_videos(
        frame_dir,
        verts_incam,
        verts_global,
        SMPLX.faces,
        SMPLX.J_regressor.to(DEVICE),
        K[0],
        fps,
        out_paths,
    )

    # results.npz for download
    flat = {}
    for key, value in results.items():
        if isinstance(value, dict):
            flat |= {f"{key}_{k}": v.numpy() for k, v in value.items()}
        else:
            flat[key] = np.asarray(value)
    np.savez(output_dir / "results.npz", **flat)
    return results, out_paths


# ---------------------------------------------------------------------------
# GPU entry points
# ---------------------------------------------------------------------------
MAX_FRAMES = 900  # 30 s at 30 fps: the demo's practical cap (quota-friendly)


DURATION_FIXED = 12.0  # per-call overhead: frame extraction + the weight stream into VRAM
DURATION_PER_FRAME = {"448x336": 0.085, "576x432": 0.107, "1024x768": 0.190}  # yolo11m
DURATION_DETECTOR_EXTRA = {"yolo11m": 0.0, "yolo11x": 0.027}  # per-frame, yolo11x is heavier
DURATION_MOVING_EXTRA = 6.0  # SimpleVO visual odometry
DURATION_MARGIN = 1.2  # ~15-25% over the worst measured run, after rounding to 5 s


def _estimate_duration(video, detector, resolution, static_cam, *args, **kwargs):
    """ZeroGPU duration estimate for one call.

    Interpolated from end-to-end timings of the live Space (900 frames = the cap):
        300 fr yolo11m 448x336 static  ->  37 s      900 fr yolo11m 448x336 static ->  88 s
        900 fr yolo11m 448x336 moving  ->  91 s      900 fr yolo11m 576x432 static -> 108 s
        900 fr yolo11x 1024x768 static -> 207 s      900 fr yolo11x 1024x768 moving -> 212 s

    We request that fit x1.2, rounded up to 5 s (each preset therefore asks ~15-25% more than
    the slowest run measured on the live Space). Keep this honest rather than generous:
    ZeroGPU compares the *requested* duration against the visitor's *remaining* quota, not
    the actual runtime, so an inflated request locks a visitor out of a call that would have
    cost them a fraction of it. (On this Blackwell runtime the `spaces` package also applies
    a `duration_factor` of 1.5, so the wall-clock reservation is 1.5x this value; the 300-frame
    default example reserves ~75 s of a free account's 300 s daily quota.)
    """
    cap = int(os.environ.get("ZERO_GPU_DURATION", "0")) or None
    frames = 300
    video = video if isinstance(video, str) else getattr(video, "name", None)
    if video:
        try:
            frames = int(V.video_lwh(video)[0])
        except Exception:
            frames = 300
    frames = max(60, min(frames, MAX_FRAMES))
    per_frame = DURATION_PER_FRAME.get(resolution, DURATION_PER_FRAME["1024x768"])
    per_frame += DURATION_DETECTOR_EXTRA.get(detector, DURATION_DETECTOR_EXTRA["yolo11x"])
    raw = DURATION_FIXED + frames * per_frame
    if not static_cam:
        raw += DURATION_MOVING_EXTRA
    seconds = int(raw * DURATION_MARGIN + 0.999)  # ceil
    seconds = ((seconds + 4) // 5) * 5  # round up to a multiple of 5 s
    if cap:
        seconds = min(seconds, cap)
    log.info(
        "reserving %ss of GPU time (%s frames, %s, %s, static_cam=%s)",
        seconds, frames, detector, resolution, static_cam,
    )
    return seconds


@spaces.GPU(duration=_estimate_duration)
def reconstruct(
    video,
    detector="yolo11m",
    resolution="448x336",
    static_cam=True,
    progress=gr.Progress(),
):
    """Reconstruct 3D body, hand and foot motion (SMPL-X) of the main person in a video.

    Args:
        video: input monocular video (mp4/webm/mov), up to ~30 s, one dominant person.
        detector: YOLO person detector preset: "yolo11m" (fast) or "yolo11x" (accurate).
        resolution: Sapiens2 crop resolution, "448x336" (balanced) to "1024x768" (accurate).
        static_cam: whether the camera is static (stationary tripod). Default True: skips
            visual odometry.

    Returns:
        A side-by-side video (input view + reconstructed mesh), and results.npz with the
        SMPL-X parameters (in-camera and global).
    """
    if video is None:
        raise gr.Error("Upload a video first")
    path = video if isinstance(video, str) else video.name
    L = V.video_lwh(path)[0]
    work = Path(tempfile.mkdtemp(prefix="exmr_"))
    if L > MAX_FRAMES:
        # trim with ffmpeg to MAX_FRAMES frames (keeps the head of the video)
        subprocess.run(
            [
                "ffmpeg", "-hide_banner", "-nostdin", "-y", "-loglevel", "error",
                "-i", path, "-frames:v", str(MAX_FRAMES), "-an",
                "-c:v", "libx264", "-preset", "veryfast", "-crf", "18",
                "-pix_fmt", "yuv420p", str(work / "input.mp4"),
            ],
            check=True,
        )
    else:
        shutil.copy(path, work / "input.mp4")
    path = str(work / "input.mp4")

    out_dir = work / "out"
    try:
        results, out_paths = _run_exmr(
            path, out_dir, detector, resolution, bool(static_cam)
        )
        side_by_side = out_paths[2]  # <name>_incam_global.mp4
        npz = out_dir / "results.npz"
        # copy the two deliverables into a stable tempdir that outlives the workdir cleanup
        keep = Path(tempfile.mkdtemp(prefix="exmr_out_"))
        shutil.copy(side_by_side, keep / "exmr_incam_global.mp4")
        shutil.copy(npz, keep / "results.npz")
        return str(keep / "exmr_incam_global.mp4"), str(keep / "results.npz")
    finally:
        shutil.rmtree(work, ignore_errors=True)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
DEFAULT_EXAMPLE = str(_HERE / "assets" / "example_dance_10s.mp4")

description = """**Reconstruct 3D body, hand and foot motion from a monocular video** with
[ExMR](https://huggingface.co/papers/2610.26213) — a transformer that refines extremities
(hands, feet) on top of whole-body tokens from SAM 3D Body, HaMeR and Sapiens2, and outputs
SMPL-X parameters.

Upload a short clip of **one person** (up to ~30 s, longer inputs are trimmed). The output
shows the input view with the reconstructed mesh overlaid, side by side with the world-frame
motion. `results.npz` holds the raw SMPL-X parameters."""

article = """### How it works
1. **Track** the main person (YOLO11 + ByteTrack) → **2D whole-body keypoints** (Sapiens2-1B).
2. **Body token** per frame from **SAM 3D Body** (DINOv3-H+); **hand tokens** from **HaMeR**.
3. Optionally estimate **camera motion** (SimpleVO) when the camera moves.
4. **ExMR** fuses everything into SMPL-X body + hand pose (with Foot/Hand Refiner modules).
5. **Render** the input-view and world-view meshes with nvdiffrast.

### Credits & licenses
Model by Tobias Weinehrbein et al. ([paper](https://huggingface.co/papers/2610.26213),
[code](https://github.com/twehrbein/ExMR)). Research-use-only components: ExMR, SMPL-X, MANO,
SAM 3D Body, Sapiens2; HaMeR is MIT; YOLO is AGPL-3.0. Example video from
[Pexels](https://www.pexels.com/video/woman-doing-dancing-exercise-in-the-garden-3690649/)
(free to use)."""

theme = gr.themes.Citrus()

with gr.Blocks(theme=theme, title="ExMR — Monocular Human Motion Capture") as demo:
    gr.Markdown("# ExMR: Extremity Motion Refinement for Monocular Human Motion Capture")
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(
                label="Input video (one person, ≤ 30 s)",
                sources=["upload"],
                format="mp4",
            )
            static_cam = gr.Checkbox(
                value=True,
                label="Static camera (tripod; uncheck if the camera moves)",
            )
            with gr.Accordion("Advanced options", open=False):
                detector = gr.Radio(
                    ["yolo11m", "yolo11x"],
                    value="yolo11m",
                    label="Person detector",
                )
                resolution = gr.Radio(
                    ["448x336", "576x432", "1024x768"],
                    value="448x336",
                    label="Sapiens2 resolution (higher = slower, slightly better)",
                )
            go = gr.Button("Reconstruct motion", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Input view + world view (side by side)")
            npz_out = gr.File(label="SMPL-X parameters (results.npz)")

    gr.Examples(
        examples=[[DEFAULT_EXAMPLE]],
        inputs=[video_in],
        fn=reconstruct,
        outputs=[video_out, npz_out],
        cache_examples=True,
        cache_mode="lazy",
        label="Example: the authors' Pexels dance clip (10 s, 300 frames)",
    )
    gr.Markdown(description)
    gr.Markdown(article)

    go.click(
        reconstruct,
        inputs=[video_in, detector, resolution, static_cam],
        outputs=[video_out, npz_out],
        api_name="reconstruct",
    )

if __name__ == "__main__":
    demo.launch(ssr_mode=False, mcp_server=True)