From e8106fee8d58c3c8071aca33122f49f43ef0dc84 Mon Sep 17 00:00:00 2001 From: Xiaohang Date: Thu, 21 May 2026 13:12:51 +0200 Subject: [PATCH 01/18] Handle PyTorch weights-only checkpoint validation --- prima/utils/weights.py | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/prima/utils/weights.py b/prima/utils/weights.py index 6234a46..e76dd1a 100644 --- a/prima/utils/weights.py +++ b/prima/utils/weights.py @@ -103,10 +103,42 @@ def _download_file( def _validate_torch_checkpoint(path: Path) -> None: + import inspect + import pickle + import zipfile + import torch + if zipfile.is_zipfile(path): + with zipfile.ZipFile(path) as checkpoint_zip: + corrupt_member = checkpoint_zip.testzip() + if corrupt_member is not None: + raise RuntimeError( + f"Checkpoint file is invalid or incomplete: {path}\n" + f"Corrupt archive member: {corrupt_member}\n" + "Please redownload the checkpoint and try again." + ) + + supports_weights_only = "weights_only" in inspect.signature(torch.load).parameters + load_kwargs = {"map_location": "cpu"} + if supports_weights_only: + load_kwargs["weights_only"] = True + try: - torch.load(path, map_location="cpu") + torch.load(path, **load_kwargs) + except pickle.UnpicklingError as exc: + message = str(exc) + if ( + supports_weights_only + and "Weights only load failed" in message + and ("Unsupported global" in message or "Unsupported class" in message) + ): + return + raise RuntimeError( + f"Checkpoint file is invalid or incomplete: {path}\n" + "Downloaded checkpoint is not loadable. " + "Please verify the uploaded Hugging Face file and try again." + ) from exc except Exception as exc: raise RuntimeError( f"Checkpoint file is invalid or incomplete: {path}\n" From 08fd52fc3a4274ab47cb3cb364a823a93fc0aa67 Mon Sep 17 00:00:00 2001 From: Xiaohang Date: Thu, 21 May 2026 13:14:21 +0200 Subject: [PATCH 02/18] add s3 inference ckpt --- README.md | 12 +----------- prima/utils/weights.py | 10 +++++----- 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 13b7360..8813183 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,3 @@ ---- -title: PRIMA Demo -emoji: 🦮 -colorFrom: blue -colorTo: green -sdk: gradio -python_version: "3.10" -app_file: app.py -startup_duration_timeout: 60m ---- - # PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation @@ -229,6 +218,7 @@ This release builds on several open-source projects, including: - [BioCLIP](https://github.com/Imageomics/BioCLIP) - [AniMer](https://github.com/luoxue-star/AniMer) - [DeepLabCut](https://github.com/DeepLabCut/DeepLabCut) +- [SAM3DB](https://github.com/facebookresearch/sam-3d-body) --- diff --git a/prima/utils/weights.py b/prima/utils/weights.py index e76dd1a..c02e593 100644 --- a/prima/utils/weights.py +++ b/prima/utils/weights.py @@ -18,7 +18,7 @@ DEFAULT_HF_REPO_ID = HF_REPO_ID DEFAULT_STAGE1_CHECKPOINT = Path("data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt") -DEFAULT_STAGE3_CHECKPOINT = Path("data/PRIMAS3/checkpoints/s3ckpt.ckpt") +DEFAULT_STAGE3_CHECKPOINT = Path("data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt") SMAL_ASSET_PATHS = [ "my_smpl_00781_4_all.pkl", @@ -29,16 +29,16 @@ STAGE1_CONFIG_ASSET_PATH = "config_s1_HYDRA.yaml" STAGE1_CHECKPOINT_ASSET_PATH = "s1ckpt_inference.ckpt" STAGE3_CONFIG_ASSET_PATH = "config_s3_HYDRA.yaml" -STAGE3_CHECKPOINT_ASSET_PATH = "s3ckpt.ckpt" +STAGE3_CHECKPOINT_ASSET_PATH = "s3ckpt_inference.ckpt" STAGE_ASSETS = { "PRIMAS1": (STAGE1_CONFIG_ASSET_PATH, STAGE1_CHECKPOINT_ASSET_PATH, "s1ckpt_inference.ckpt"), - "PRIMAS3": (STAGE3_CONFIG_ASSET_PATH, STAGE3_CHECKPOINT_ASSET_PATH, "s3ckpt.ckpt"), + "PRIMAS3": (STAGE3_CONFIG_ASSET_PATH, STAGE3_CHECKPOINT_ASSET_PATH, "s3ckpt_inference.ckpt"), } STAGE_CHECKPOINTS = { "PRIMAS1": Path("PRIMAS1/checkpoints/s1ckpt_inference.ckpt"), - "PRIMAS3": Path("PRIMAS3/checkpoints/s3ckpt.ckpt"), + "PRIMAS3": Path("PRIMAS3/checkpoints/s3ckpt_inference.ckpt"), } PathLike = Union[str, Path] @@ -274,7 +274,7 @@ def _ensure_assets_for_checkpoint( f" config: {config_path}\n" "Auto-download supports the standard PRIMA demo layouts only:\n" " data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt\n" - " data/PRIMAS3/checkpoints/s3ckpt.ckpt\n" + " data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt\n" "Pass one of those paths, or download/copy your custom checkpoint manually." ) From 71c1e2c56551c50be496435c851deb8b6e4e6ad0 Mon Sep 17 00:00:00 2001 From: Xiaohang Date: Thu, 21 May 2026 13:38:21 +0200 Subject: [PATCH 03/18] add egl and osmesa for renderer --- prima/utils/mesh_renderer.py | 5 +++-- prima/utils/renderer.py | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/prima/utils/mesh_renderer.py b/prima/utils/mesh_renderer.py index 707443b..522bc1a 100644 --- a/prima/utils/mesh_renderer.py +++ b/prima/utils/mesh_renderer.py @@ -8,10 +8,11 @@ """ import os +from ctypes.util import find_library if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin': - # Use software OpenGL in headless Linux environments (e.g., Hugging Face Spaces). - os.environ['PYOPENGL_PLATFORM'] = 'osmesa' + # Prefer OSMesa; fall back to EGL where available. + os.environ['PYOPENGL_PLATFORM'] = 'osmesa' if find_library('OSMesa') else 'egl' import torch from torchvision.utils import make_grid import numpy as np diff --git a/prima/utils/renderer.py b/prima/utils/renderer.py index 989052b..b8f4e4a 100644 --- a/prima/utils/renderer.py +++ b/prima/utils/renderer.py @@ -10,10 +10,11 @@ import os +from ctypes.util import find_library if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin': - # Use software OpenGL in headless Linux environments (e.g., Hugging Face Spaces). - os.environ['PYOPENGL_PLATFORM'] = 'osmesa' + # Prefer OSMesa; fall back to EGL where available. + os.environ['PYOPENGL_PLATFORM'] = 'osmesa' if find_library('OSMesa') else 'egl' import torch import numpy as np import pyrender @@ -430,4 +431,3 @@ def add_point_lighting(self, scene, cam_node, color=np.ones(3), intensity=1.0): scene.add_node(node) - From 92493159b7fc10edcafce521a76241ecbe2c9c50 Mon Sep 17 00:00:00 2001 From: Xiaohang Date: Thu, 21 May 2026 13:54:20 +0200 Subject: [PATCH 04/18] move to gpu device --- eval.py | 1 + 1 file changed, 1 insertion(+) diff --git a/eval.py b/eval.py index b7031c5..659b7eb 100644 --- a/eval.py +++ b/eval.py @@ -25,6 +25,7 @@ def main(args): default_cfg = get_config(args.default_eval_config) model = PRIMA.load_from_checkpoint(args.checkpoint, cfg=cfg, strict=False) model.eval() + model = model.to(args.device) smal_evaluator = Evaluator(smal_model=model.smal, image_size=cfg.MODEL.IMAGE_SIZE) cfg_eval_dataset = dict(default_cfg.DATASETS) From c3a3e5ff0a4eab4a5ebecbd208aaae2dc4f65fcf Mon Sep 17 00:00:00 2001 From: Xiaohang Date: Thu, 21 May 2026 13:56:52 +0200 Subject: [PATCH 05/18] update s3 inference ckpt --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8813183..bb44f5a 100644 --- a/README.md +++ b/README.md @@ -110,7 +110,7 @@ Expected files in that Hugging Face repo root: Optional Stage 3 prefetch expects: - `config_s3_HYDRA.yaml` -- `s3ckpt.ckpt` +- `s3ckpt_inference.ckpt` ### Demo (without TTA) From 60dcea8bd34a5c502eb22f2f34d434f0a5063615 Mon Sep 17 00:00:00 2001 From: Xiaohang Date: Thu, 21 May 2026 14:19:43 +0200 Subject: [PATCH 06/18] update pyproject' --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d76bd36..dc163e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "prima-animal" version = "0.1.7" description = "PRIMA: 3D animal pose and shape estimation" readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.10" authors = [ { name = "Xiaohang Yu", email = "xiaohang.yu@epfl.ch" }, From 9bee30d67258097729c0dfd4588c69e43a75795a Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 21 May 2026 15:12:00 +0200 Subject: [PATCH 07/18] fix comment for example checkpoint path in demo_tta.sh --- demo_tta.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demo_tta.sh b/demo_tta.sh index 8f0f111..58eef2d 100644 --- a/demo_tta.sh +++ b/demo_tta.sh @@ -4,7 +4,7 @@ # # This standard path is auto-downloaded from the PRIMA Hugging Face repo if missing. # To use another local checkpoint instead, update this path. -# For example: checkpoint='data/PRIMAS3/checkpoints/s3ckpt.ckpt' +# For example: checkpoint='data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt' checkpoint='data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt' python3 demo_tta.py \ From 39bde99650733cb12021052e3f10c0c9b418aef4 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 21 May 2026 15:12:18 +0200 Subject: [PATCH 08/18] fix example checkpoint path in demo.sh --- demo.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demo.sh b/demo.sh index e45cc37..6dace89 100644 --- a/demo.sh +++ b/demo.sh @@ -3,7 +3,7 @@ # # If this local file is missing, it will be downloaded from the PRIMA Hugging Face repo. # To use another local checkpoint instead, update this path. -# For example: checkpoint='data/PRIMAS3/checkpoints/s3ckpt.ckpt' +# For example: checkpoint='data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt' checkpoint='data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt' python demo.py \ From 5a449715e0e360edaeaffa05a8d978d0683920c6 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 21 May 2026 19:18:32 +0200 Subject: [PATCH 09/18] update OpenGL platform preference to use EGL with surfaceless option --- prima/utils/mesh_renderer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/prima/utils/mesh_renderer.py b/prima/utils/mesh_renderer.py index 522bc1a..7501004 100644 --- a/prima/utils/mesh_renderer.py +++ b/prima/utils/mesh_renderer.py @@ -11,8 +11,10 @@ from ctypes.util import find_library if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin': - # Prefer OSMesa; fall back to EGL where available. - os.environ['PYOPENGL_PLATFORM'] = 'osmesa' if find_library('OSMesa') else 'egl' + # Prefer EGL; PyOpenGL's OSMesa bindings can lack symbols required by pyrender. + os.environ['PYOPENGL_PLATFORM'] = 'egl' if find_library('EGL') else 'osmesa' + if os.environ['PYOPENGL_PLATFORM'] == 'egl': + os.environ.setdefault('EGL_PLATFORM', 'surfaceless') import torch from torchvision.utils import make_grid import numpy as np From ffb58f4a4039cc73dd4559e69a088c2ffa69c7fe Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 21 May 2026 19:18:44 +0200 Subject: [PATCH 10/18] update OpenGL platform preference to use EGL with surfaceless option --- prima/utils/renderer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/prima/utils/renderer.py b/prima/utils/renderer.py index b8f4e4a..71ae021 100644 --- a/prima/utils/renderer.py +++ b/prima/utils/renderer.py @@ -13,8 +13,10 @@ from ctypes.util import find_library if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin': - # Prefer OSMesa; fall back to EGL where available. - os.environ['PYOPENGL_PLATFORM'] = 'osmesa' if find_library('OSMesa') else 'egl' + # Prefer EGL; PyOpenGL's OSMesa bindings can lack symbols required by pyrender. + os.environ['PYOPENGL_PLATFORM'] = 'egl' if find_library('EGL') else 'osmesa' + if os.environ['PYOPENGL_PLATFORM'] == 'egl': + os.environ.setdefault('EGL_PLATFORM', 'surfaceless') import torch import numpy as np import pyrender @@ -429,5 +431,3 @@ def add_point_lighting(self, scene, cam_node, color=np.ones(3), intensity=1.0): if scene.has_node(node): continue scene.add_node(node) - - From fd920ae724fc9e2f873c240ca84f84928eecf599 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 21 May 2026 19:19:15 +0200 Subject: [PATCH 11/18] refactor: enhance progress reporting in animal detection and TTA process --- app.py | 58 +++++++++++++++++++++++++++++++++++++++++++++------------- 1 file changed, 45 insertions(+), 13 deletions(-) diff --git a/app.py b/app.py index 2d0216c..f2cfbef 100644 --- a/app.py +++ b/app.py @@ -24,12 +24,13 @@ import argparse import concurrent.futures import os +import queue import sys import tempfile import time import traceback from types import SimpleNamespace -from typing import List, Tuple +from typing import Callable, List, Tuple from pathlib import Path import cv2 @@ -77,11 +78,11 @@ def _gradio_examples_for_interface() -> List[List]: return [] rows: List[List] = [] template: List[Tuple[str, float, int, float, float, bool, bool]] = [ - ("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, True), - ("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, True), - ("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, True), - ("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, True), - ("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, True), + ("demo_data/000000015956_horse.png", 1e-6, 0, 0.7, 0.1, False, True), + ("demo_data/n02412080_12159.png", 1e-6, 0, 0.7, 0.1, False, True), + ("demo_data/000000315905_zebra.jpg", 1e-6, 0, 0.7, 0.1, False, True), + ("demo_data/beagle.jpg", 1e-6, 0, 0.7, 0.1, False, True), + ("demo_data/shepherd_hati.jpg", 1e-6, 0, 0.7, 0.1, False, True), ] for rel, *rest in template: p = _REPO_ROOT / rel @@ -209,6 +210,7 @@ def _collect_animal_results( kp_conf_thresh: float, side_view: bool, save_mesh: bool, + progress_callback: Callable[[str], None] | None = None, ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], str | None, str | None]: """Run detection + PRIMA + SuperAnimal + TTA on a single RGB image. @@ -229,10 +231,16 @@ def _collect_animal_results( tta_optimize, ) + def report(message: str) -> None: + if progress_callback is not None: + progress_callback(message) + if int(tta_num_iters) > 0 and not SUPER_ANIMAL_ARGS.saved_2d_model_path: + report("Resolving SuperAnimal weights...") SUPER_ANIMAL_ARGS.saved_2d_model_path = resolve_sa_weights_path("") # Detect animals + report("Detecting animals...") img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) if detector is None: # Fallback for environments where Detectron2 is unavailable: process full image as one crop. @@ -248,6 +256,7 @@ def _collect_animal_results( if len(boxes) == 0: return [], [], [], None, None + report(f"Detected {len(boxes)} animal(s). Preparing crops...") dataset = ViTDetDataset(model_cfg, img_bgr, boxes) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0) @@ -259,9 +268,11 @@ def _collect_animal_results( img_token = next(tempfile._get_candidate_names()) - for batch in dataloader: + total_batches = len(dataloader) + for batch_idx, batch in enumerate(dataloader, start=1): batch = recursive_to(batch, device) + report(f"Animal {batch_idx}/{total_batches}: running PRIMA...") with torch.no_grad(): out_before = model(batch) @@ -271,6 +282,7 @@ def _collect_animal_results( img_fn = f"{img_token}" from demo_tta import render_and_save # imported lazily to avoid circular issues + report(f"Animal {batch_idx}/{total_batches}: rendering before TTA...") render_and_save( renderer, cam_crop_to_full_fn, @@ -296,6 +308,7 @@ def _collect_animal_results( before_mesh_paths.append(before_obj_path) if int(tta_num_iters) <= 0: + report(f"Animal {batch_idx}/{total_batches}: rendering final output...") render_and_save( renderer, cam_crop_to_full_fn, @@ -322,6 +335,7 @@ def _collect_animal_results( continue # Prepare patch for SuperAnimal + report(f"Animal {batch_idx}/{total_batches}: running SuperAnimal keypoints...") patch_rgb = denorm_patch_to_rgb(batch["img"][0]) with tempfile.TemporaryDirectory(prefix=f"dlc_{img_fn}_{animal_id}_") as tmp_dir: bodyparts_xyc = run_superanimal_on_patch(patch_rgb, SUPER_ANIMAL_ARGS, tmp_dir) @@ -352,6 +366,7 @@ def _collect_animal_results( gt_kpts_norm = torch.from_numpy(kpts_norm[None]).to(device=device, dtype=batch["img"].dtype) # Run TTA + report(f"Animal {batch_idx}/{total_batches}: running TTA ({int(tta_num_iters)} iterations)...") out_after = tta_optimize( model, batch, @@ -360,6 +375,7 @@ def _collect_animal_results( lr=float(tta_lr), ) + report(f"Animal {batch_idx}/{total_batches}: rendering after TTA...") render_and_save( renderer, cam_crop_to_full_fn, @@ -387,6 +403,7 @@ def _collect_animal_results( first_before_mesh = before_mesh_paths[0] if before_mesh_paths else None first_after_mesh = after_mesh_paths[0] if after_mesh_paths else None + report("Collecting outputs...") return before_imgs, after_imgs, kpt_imgs, first_before_mesh, first_after_mesh @@ -487,6 +504,11 @@ def gradio_inference( save_mesh=save_mesh, ) else: + stage_updates: queue.Queue[str] = queue.Queue() + + def report_stage(message: str) -> None: + stage_updates.put(message) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: fut = pool.submit( _collect_animal_results, @@ -504,20 +526,30 @@ def gradio_inference( kp_conf_thresh, side_view, save_mesh, + report_stage, ) t0 = time.monotonic() + latest_stage = "Starting inference..." while True: + while True: + try: + latest_stage = stage_updates.get_nowait() + except queue.Empty: + break + else: + elapsed = int(time.monotonic() - t0) + yield None, None, None, f"{latest_stage}\nElapsed: {elapsed}s" try: before_imgs, after_imgs, kpt_imgs, mesh_before, mesh_after = fut.result( - timeout=hb + timeout=1.0 ) break except concurrent.futures.TimeoutError: elapsed = int(time.monotonic() - t0) yield None, None, None, ( - f"Inference still running ({elapsed}s). " - f"Detection, SuperAnimal, and TTA can take several minutes; " - f"updates every ~{int(hb)}s keep the connection alive." + f"{latest_stage}\n" + f"Elapsed: {elapsed}s\n" + "CPU inference can take several minutes." ) except Exception: yield None, None, None, f"Inference failed:\n{traceback.format_exc()}" @@ -558,7 +590,7 @@ def gradio_inference( label="TTA iterations", minimum=0, maximum=100, - value=30, + value=0, step=1, info="Set to 0 to disable TTA and reuse the initial PRIMA prediction.", ), @@ -624,4 +656,4 @@ def parse_args() -> argparse.Namespace: if _should_preload_assets(): _preload_assets_once(args.checkpoint) demo = build_demo(checkpoint_path=args.checkpoint, out_folder=args.out_folder) - demo.launch(inbrowser=False) + demo.launch(inbrowser=False, ssr_mode=False) From 91b7910359da9ba7286043be404c620ca510b6b9 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 21 May 2026 19:19:28 +0200 Subject: [PATCH 12/18] refactor: update deployment script to sync only tracked files and enhance error handling --- scripts/deploy_hf_space.sh | 96 +++++++++++++++++++++++++++++++------- 1 file changed, 78 insertions(+), 18 deletions(-) diff --git a/scripts/deploy_hf_space.sh b/scripts/deploy_hf_space.sh index aeed09a..f31ca1d 100755 --- a/scripts/deploy_hf_space.sh +++ b/scripts/deploy_hf_space.sh @@ -2,8 +2,10 @@ # Deploy working tree to Hugging Face Space MLAdaptiveIntelligence/PRIMA-demo. # # Demo PNG/JPG are tracked with Git LFS (Hugging Face Hub Xet bridge); see .gitattributes. -# We rsync the working tree (not ``git archive``) so LFS-tracked files are real bytes here, -# then ``git add`` stores them as LFS objects on push. +# We rsync only the Git-tracked files needed by app.py from the working tree +# (not ``git archive``), so tracked LFS files are materialized bytes while +# untracked local files and non-Space project files stay out. Then ``git add`` +# stores matching files as LFS objects on push. # # Prerequisites: brew install git-lfs git-xet && git xet install && git lfs install set -euo pipefail @@ -23,20 +25,70 @@ TMP="$(mktemp -d)" cleanup() { rm -rf "$TMP"; } trap cleanup EXIT -echo "[deploy] Rsync working tree from ${ROOT} ..." -rsync -a \ - --exclude=".git/" \ - --exclude="__pycache__/" \ - --exclude="*.pyc" \ - --exclude=".DS_Store" \ - --exclude="data/" \ - --exclude=".venv/" \ - --exclude="venv/" \ - --exclude=".pytest_cache/" \ - --exclude=".tmp_gradio_info.json" \ - --exclude="demo_out_tta_gradio/" \ - --exclude=".gradio/" \ - "${ROOT}/" "${TMP}/" +SPACE_SYNC_PATHS=( + ".gitattributes" + "README.md" + "requirements.txt" + "pyproject.toml" + "app.py" + "demo_tta.py" + "chumpy" + "configs/sa_finetune_hrnet_w32.yaml" + "demo_data" + "images/teaser.png" + "prima" +) +SPACE_EXTRA_FILES=( + "packages.txt" +) + +echo "[deploy] Rsync Git-tracked Space files from ${ROOT} ..." +printf '[deploy] %s\n' "${SPACE_SYNC_PATHS[@]}" +missing_tracked=() +for path in "${SPACE_SYNC_PATHS[@]}"; do + if [[ -z "$(git ls-files -- "$path")" ]]; then + missing_tracked+=("$path") + fi +done +if [[ "${#missing_tracked[@]}" -gt 0 ]]; then + printf '[deploy] ERROR: Space sync path is not tracked by Git: %s\n' "${missing_tracked[@]}" >&2 + echo "[deploy] Add required new files with git add, or remove them from SPACE_SYNC_PATHS." >&2 + exit 1 +fi +git ls-files -z -- "${SPACE_SYNC_PATHS[@]}" | rsync -a --from0 --files-from=- "${ROOT}/" "${TMP}/" + +echo "[deploy] Rsync explicit Space config files from ${ROOT} ..." +for path in "${SPACE_EXTRA_FILES[@]}"; do + if [[ ! -f "$path" ]]; then + echo "[deploy] ERROR: Missing required Space config file: $path" >&2 + exit 1 + fi + printf '[deploy] %s\n' "$path" + rsync -a --relative "$path" "$TMP/" +done + +README_FILE="${TMP}/README.md" +if ! sed -n '1,20p' "$README_FILE" | grep -q '^sdk: gradio$'; then + echo "[deploy] Adding Hugging Face Space YAML front matter to README.md ..." + README_TMP="${README_FILE}.tmp" + { + cat <<'YAML' +--- +title: PRIMA Demo +emoji: 🦮 +colorFrom: blue +colorTo: green +sdk: gradio +python_version: "3.10" +app_file: app.py +startup_duration_timeout: 60m +--- + +YAML + cat "$README_FILE" + } > "$README_TMP" + mv "$README_TMP" "$README_FILE" +fi cd "$TMP" @@ -46,7 +98,15 @@ git lfs install git add -A git -c user.email="space-deploy@users.noreply.github.com" -c user.name="HF Space deploy" commit -q -m "Deploy snapshot (LFS for demo images per .gitattributes)" -git remote add hf "$SPACE_URL" +PUSH_URL="$SPACE_URL" +if [[ "$PUSH_URL" == https://huggingface.co/* && -z "${HF_TOKEN:-}" && -f "${HF_HOME:-$HOME/.cache/huggingface}/token" ]]; then + HF_TOKEN="$(<"${HF_HOME:-$HOME/.cache/huggingface}/token")" +fi +if [[ "$PUSH_URL" == https://huggingface.co/* && -n "${HF_TOKEN:-}" ]]; then + PUSH_URL="${PUSH_URL/https:\/\/huggingface.co/https:\/\/hf_user:${HF_TOKEN}@huggingface.co}" +fi + +git remote add hf "$PUSH_URL" echo "[deploy] Force-pushing to Hugging Face Space ..." -GIT_TERMINAL_PROMPT=0 git -c credential.helper=osxkeychain push hf HEAD:main --force +GIT_TERMINAL_PROMPT=0 git -c credential.helper= push hf HEAD:main --force echo "[deploy] Done." From b8ccb0d3589ba67e1126d6e47555c8c8287b864e Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 28 May 2026 17:08:49 +0200 Subject: [PATCH 13/18] refactor: improve reporting for animal detection process in _collect_animal_results function --- app.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/app.py b/app.py index f2cfbef..d41b252 100644 --- a/app.py +++ b/app.py @@ -240,13 +240,14 @@ def report(message: str) -> None: SUPER_ANIMAL_ARGS.saved_2d_model_path = resolve_sa_weights_path("") # Detect animals - report("Detecting animals...") img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) if detector is None: # Fallback for environments where Detectron2 is unavailable: process full image as one crop. + report("Detectron2 unavailable; using full-image crop...") h, w = img_bgr.shape[:2] boxes = np.array([[0.0, 0.0, float(max(1, w - 1)), float(max(1, h - 1))]], dtype=np.float32) else: + report("Detecting animals with Detectron2...") det_out = detector(img_bgr) det_instances = det_out["instances"] From b753959b66627d562a158cdd40ed459a18e06638 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 28 May 2026 17:47:58 +0200 Subject: [PATCH 14/18] Resolve ti_dev merge follow-ups --- README.md | 11 +++- app.py | 104 ++++++++++++++++++++--------------- packages.txt | 7 +++ prima/utils/mesh_renderer.py | 6 +- prima/utils/renderer.py | 8 +-- scripts/local_infer.py | 1 - 6 files changed, 83 insertions(+), 54 deletions(-) create mode 100644 packages.txt diff --git a/README.md b/README.md index 22ad91f..1d20487 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,11 @@ Options: - `PRIMA_VENV=.venv ./scripts/clean_install_local.sh --skip-data` — skip the large `setup_demo_data` download if `data/` is already populated. - `./scripts/clean_install_local.sh --wipe-data --force-data` — delete downloaded `data/` assets and redownload. - `./scripts/clean_install_local.sh --no-editable` — only `requirements.txt` (no `pip install -e .`); use if editable install fails and you will install the training stack via conda as in the PyPI section above. You still need **Python 3.10+** for Gradio 5.1+. The smoke test sets `PYTHONPATH` to the repo root so `import prima` works without an editable install. -- **`requirements.txt` pins `deeplabcut==3.0.0rc14`** (SuperAnimal PyTorch API). On macOS, `clean_install_local.sh` installs a PyTables wheel first, then DLC 3.x. Full check: `./scripts/test_local_full.sh`. +- **macOS / DeepLabCut:** `requirements.txt` pins `deeplabcut==3.0.0rc14` + for the SuperAnimal PyTorch API. On macOS, `clean_install_local.sh` installs + it separately after a compatible PyTables wheel (`tables>=3.9.2,<3.11`) to + avoid Apple Silicon build issues. Validate the local setup with + `./scripts/test_local_full.sh`. After `requirements.txt`, the script runs **`pip install --no-deps -e .`** so the `prima` package is registered without re-resolving `pyproject.toml` (which would pull **Detectron2** from git again). Install Detectron2 separately if needed: `pip install 'git+https://github.com/facebookresearch/detectron2.git'`. @@ -186,8 +190,9 @@ Then from a clean checkout with LFS files present, redeploy the Space (same as ` ./scripts/clean_redeploy_hf_space.sh ``` -The script rsyncs the working tree (not `git archive`) so image files are materialized -before `git add` turns them into LFS blobs. +The script rsyncs only the Git-tracked files needed by the Space from the +working tree (not `git archive`) so image files are materialized before +`git add` turns them into LFS blobs. --- diff --git a/app.py b/app.py index b388d51..ff19d0c 100644 --- a/app.py +++ b/app.py @@ -22,15 +22,17 @@ """ import argparse +import concurrent.futures import os import queue import sys import tempfile +import time import traceback from dataclasses import dataclass from functools import lru_cache from types import SimpleNamespace -from typing import Callable, List, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple from pathlib import Path # macOS: PyRender (OpenGL) and DeepLabCut/pyglet must run on the main thread. @@ -189,14 +191,7 @@ def _gradio_examples_for_interface(profile: DemoProfile) -> List[List]: if _is_truthy_env("PRIMA_DISABLE_GRADIO_EXAMPLES"): return [] rows: List[List] = [] - template: List[Tuple[str, float, int, float, float, bool, bool]] = [ - ("demo_data/000000015956_horse.png", 1e-6, 0, 0.7, 0.1, False, True), - ("demo_data/n02412080_12159.png", 1e-6, 0, 0.7, 0.1, False, True), - ("demo_data/000000315905_zebra.jpg", 1e-6, 0, 0.7, 0.1, False, True), - ("demo_data/beagle.jpg", 1e-6, 0, 0.7, 0.1, False, True), - ("demo_data/shepherd_hati.jpg", 1e-6, 0, 0.7, 0.1, False, True), - ] - for rel, *rest in template: + for rel, *rest in profile.example_rows: p = _REPO_ROOT / rel if p.is_file(): rows.append([str(p), *rest]) @@ -371,7 +366,8 @@ def _collect_animal_results( kp_conf_thresh: float, side_view: bool, save_mesh: bool, - progress_callback: Callable[[str], None] | None = None, + boxes: Optional[np.ndarray] = None, + progress_callback: Optional[Callable[[str], None]] = None, ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], str | None, str | None]: """Run detection + PRIMA + SuperAnimal + TTA on a single RGB image. @@ -401,21 +397,14 @@ def report(message: str) -> None: SUPER_ANIMAL_ARGS.saved_2d_model_path = resolve_sa_weights_path("") img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) - if detector is None: - # Fallback for environments where Detectron2 is unavailable: process full image as one crop. - report("Detectron2 unavailable; using full-image crop...") - h, w = img_bgr.shape[:2] - boxes = np.array([[0.0, 0.0, float(max(1, w - 1)), float(max(1, h - 1))]], dtype=np.float32) - else: - report("Detecting animals with Detectron2...") - det_out = detector(img_bgr) - det_instances = det_out["instances"] - - boxes, suppressed = select_animal_boxes(det_instances, score_threshold=float(det_thresh)) - if suppressed > 0: - print(f"[INFO] Suppressed {suppressed} duplicate animal detection(s)") - if len(boxes) == 0: - return [], [], [], None, None + if boxes is None: + if detector is None: + report("Detectron2 unavailable; using full-image crop...") + else: + report("Detecting animals with Detectron2...") + boxes = _detect_animal_boxes(detector, img_bgr, det_thresh) + if boxes is None: + return [], [], [], None, None report(f"Detected {len(boxes)} animal(s). Preparing crops...") dataset = ViTDetDataset(model_cfg, img_bgr, boxes) @@ -662,7 +651,35 @@ def gradio_inference( None, "No animal detected. Try lowering the detection threshold or another image.", ) - else: + return + yield ( + None, + None, + None, + f"Detected {len(boxes)} animal region(s). Running PRIMA (+ SuperAnimal/TTA if enabled)...", + ) + + def run_collect(progress_callback: Optional[Callable[[str], None]] = None): + return _collect_animal_results( + runtime_cache["model"], + runtime_cache["model_cfg"], + runtime_cache["renderer"], + runtime_cache["cam_crop_to_full_fn"], + runtime_cache["device"], + runtime_cache["detector"], + out_folder, + img_rgb, + tta_lr=tta_lr, + tta_num_iters=tta_num_iters, + det_thresh=det_thresh, + kp_conf_thresh=kp_conf_thresh, + side_view=side_view, + save_mesh=save_mesh, + boxes=boxes, + progress_callback=progress_callback, + ) + + if _should_use_gradio_queue(profile): stage_updates: queue.Queue[str] = queue.Queue() def report_stage(message: str) -> None: @@ -670,21 +687,7 @@ def report_stage(message: str) -> None: with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: fut = pool.submit( - _collect_animal_results, - runtime_cache["model"], - runtime_cache["model_cfg"], - runtime_cache["renderer"], - runtime_cache["cam_crop_to_full_fn"], - runtime_cache["device"], - runtime_cache["detector"], - out_folder, - img_rgb, - tta_lr, - tta_num_iters, - det_thresh, - kp_conf_thresh, - side_view, - save_mesh, + run_collect, report_stage, ) t0 = time.monotonic() @@ -710,6 +713,8 @@ def report_stage(message: str) -> None: f"Elapsed: {elapsed}s\n" "CPU inference can take several minutes." ) + else: + before_imgs, after_imgs, kpt_imgs, mesh_before, mesh_after = run_collect() except Exception: yield None, None, None, f"Inference failed:\n{traceback.format_exc()}" return @@ -748,8 +753,8 @@ def report_stage(message: str) -> None: gr.Slider( label="TTA iterations", minimum=0, - maximum=100, - value=0, + maximum=profile.max_tta_iters, + value=profile.default_tta_iters, step=1, info="Set to 0 to disable TTA and reuse the initial PRIMA prediction.", ), @@ -809,5 +814,16 @@ def parse_args() -> argparse.Namespace: profile = get_demo_profile() if _should_preload_assets(profile): _preload_assets_once(args.checkpoint) - demo = build_demo(checkpoint_path=args.checkpoint, out_folder=args.out_folder) + runtime_cache: Optional[Dict[str, Any]] = None + if ( + sys.platform == "darwin" + and profile.mode == "local" + and _is_truthy_env("PRIMA_WARMUP") + ): + runtime_cache = _warmup_runtime_cache(args.checkpoint, profile) + demo = build_demo( + checkpoint_path=args.checkpoint, + out_folder=args.out_folder, + runtime_cache=runtime_cache, + ) demo.launch(inbrowser=False, ssr_mode=False) diff --git a/packages.txt b/packages.txt new file mode 100644 index 0000000..aaca018 --- /dev/null +++ b/packages.txt @@ -0,0 +1,7 @@ +libosmesa6 +libgl1 +libgl1-mesa-dri +libegl-mesa0 +libegl1 +libglx-mesa0 +libgles2 diff --git a/prima/utils/mesh_renderer.py b/prima/utils/mesh_renderer.py index 522bc1a..7501004 100644 --- a/prima/utils/mesh_renderer.py +++ b/prima/utils/mesh_renderer.py @@ -11,8 +11,10 @@ from ctypes.util import find_library if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin': - # Prefer OSMesa; fall back to EGL where available. - os.environ['PYOPENGL_PLATFORM'] = 'osmesa' if find_library('OSMesa') else 'egl' + # Prefer EGL; PyOpenGL's OSMesa bindings can lack symbols required by pyrender. + os.environ['PYOPENGL_PLATFORM'] = 'egl' if find_library('EGL') else 'osmesa' + if os.environ['PYOPENGL_PLATFORM'] == 'egl': + os.environ.setdefault('EGL_PLATFORM', 'surfaceless') import torch from torchvision.utils import make_grid import numpy as np diff --git a/prima/utils/renderer.py b/prima/utils/renderer.py index b05a0bf..f20e41f 100644 --- a/prima/utils/renderer.py +++ b/prima/utils/renderer.py @@ -13,8 +13,10 @@ from ctypes.util import find_library if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin': - # Prefer OSMesa; fall back to EGL where available. - os.environ['PYOPENGL_PLATFORM'] = 'osmesa' if find_library('OSMesa') else 'egl' + # Prefer EGL; PyOpenGL's OSMesa bindings can lack symbols required by pyrender. + os.environ['PYOPENGL_PLATFORM'] = 'egl' if find_library('EGL') else 'osmesa' + if os.environ['PYOPENGL_PLATFORM'] == 'egl': + os.environ.setdefault('EGL_PLATFORM', 'surfaceless') import torch import numpy as np import pyrender @@ -438,5 +440,3 @@ def add_point_lighting(self, scene, cam_node, color=np.ones(3), intensity=1.0): if scene.has_node(node): continue scene.add_node(node) - - diff --git a/scripts/local_infer.py b/scripts/local_infer.py index 1de893b..2ed7081 100755 --- a/scripts/local_infer.py +++ b/scripts/local_infer.py @@ -109,4 +109,3 @@ def main() -> int: if __name__ == "__main__": raise SystemExit(main()) - From ca27a041eef9eec0c289c5de26953750c187e8d4 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 28 May 2026 18:09:24 +0200 Subject: [PATCH 15/18] fix(deploy): remove Detectron2 from Space requirements for fallback to full-image crops --- scripts/deploy_hf_space.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/scripts/deploy_hf_space.sh b/scripts/deploy_hf_space.sh index f31ca1d..b45909c 100755 --- a/scripts/deploy_hf_space.sh +++ b/scripts/deploy_hf_space.sh @@ -68,6 +68,12 @@ for path in "${SPACE_EXTRA_FILES[@]}"; do done README_FILE="${TMP}/README.md" +REQ_FILE="${TMP}/requirements.txt" + +echo "[deploy] Removing Detectron2 from Space requirements (app falls back to full-image crops) ..." +grep -vE '^[[:space:]]*detectron2([[:space:]]|@|$)' "$REQ_FILE" > "${REQ_FILE}.tmp" +mv "${REQ_FILE}.tmp" "$REQ_FILE" + if ! sed -n '1,20p' "$README_FILE" | grep -q '^sdk: gradio$'; then echo "[deploy] Adding Hugging Face Space YAML front matter to README.md ..." README_TMP="${README_FILE}.tmp" From 4a52c35e2969f38493ec7ace5a63ddb0319473c1 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 28 May 2026 18:09:43 +0200 Subject: [PATCH 16/18] fix(demo): update description to clarify fallback to full-image crop in Hugging Face Space demo --- app.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/app.py b/app.py index ff19d0c..5587531 100644 --- a/app.py +++ b/app.py @@ -159,9 +159,9 @@ def resolve_detectron_device(self) -> str: ("demo_data/000000315905_zebra.jpg", 1e-6, 0, 0.7, 0.1, False, False), ), description=( - "**Hugging Face Space (cpu-basic)** — lightweight demo: **CPU-only**, Detectron2 **R50-FPN**, " - "PRIMA inference. TTA is optional (0 by default; increases runtime). Mesh `.obj` export is off " - "by default to save time and disk." + "**Hugging Face Space (cpu-basic)** — lightweight demo: **CPU-only** PRIMA inference. " + "The Space build skips Detectron2 and uses a full-image crop fallback. TTA is optional " + "(0 by default; increases runtime). Mesh `.obj` export is off by default to save time and disk." ), interface_title="PRIMA on Hugging Face — lightweight CPU demo", ) From c40d92224c0451321f1a1d158903c77328b6ff34 Mon Sep 17 00:00:00 2001 From: ti Date: Thu, 28 May 2026 18:20:13 +0200 Subject: [PATCH 17/18] fix(readme): update detector information and clarify fallback mechanism for Hugging Face Space --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 1d20487..46f7906 100644 --- a/README.md +++ b/README.md @@ -163,7 +163,7 @@ The `s1ckpt_inference.ckpt` checkpoint is downloaded automatically if missing. | | **Local** (`python app.py`) | **Hugging Face Space** | |--|--|--| | PRIMA device | GPU if available, else CPU | CPU only | -| Detectron2 | X-101-FPN | R50-FPN (lighter) | +| Detector | Detectron2 X-101-FPN | full-image crop fallback | | Default TTA iterations | 30 | 0 (PRIMA-only by default) | | Save `.obj` meshes | on | off | | Preload checkpoint at startup | off | on | @@ -193,6 +193,8 @@ Then from a clean checkout with LFS files present, redeploy the Space (same as ` The script rsyncs only the Git-tracked files needed by the Space from the working tree (not `git archive`) so image files are materialized before `git add` turns them into LFS blobs. +During deployment, `detectron2` is removed from the Space `requirements.txt`; +the app uses its full-image crop fallback on the CPU Space. --- From e3bc151c329239f5045cc2e1626fb46266ae503e Mon Sep 17 00:00:00 2001 From: ti Date: Fri, 29 May 2026 10:39:37 +0200 Subject: [PATCH 18/18] feat(demo): add server name and port configuration for Gradio interface --- app.py | 38 ++++++++++++++++++++++++++++++-------- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/app.py b/app.py index 5587531..39a432d 100644 --- a/app.py +++ b/app.py @@ -68,6 +68,8 @@ # Output folder for rendered images/meshes and keypoints DEFAULT_OUT_FOLDER = "demo_out_tta_gradio" +DEFAULT_SERVER_NAME = os.environ.get("PRIMA_GRADIO_HOST", "0.0.0.0") +DEFAULT_SERVER_PORT = int(os.environ.get("PRIMA_GRADIO_PORT", "7860")) _D2_R50_CFG = "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml" _D2_R50_URL = ( @@ -128,8 +130,8 @@ def resolve_detectron_device(self) -> str: ("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, True), ("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, True), ("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, True), - ("demo_data/beagle.jpg", 1e-6, 0, 0.7, 0.1, False, True), - ("demo_data/shepherd_hati.jpg", 1e-6, 0, 0.7, 0.1, False, True), + ("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, True), + ("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, True), ), description=( "**Local demo** — full pipeline on your machine (GPU when available).\n\n" @@ -148,20 +150,23 @@ def resolve_detectron_device(self) -> str: detectron_config_yaml=_D2_R50_CFG, detectron_weights_url=_D2_R50_URL, detectron_device="cpu", - default_tta_iters=0, + default_tta_iters=30, max_tta_iters=30, default_save_mesh=False, default_side_view=False, preload_assets=True, example_rows=( - ("demo_data/beagle.jpg", 1e-6, 0, 0.7, 0.1, False, False), - ("demo_data/000000015956_horse.png", 1e-6, 0, 0.7, 0.1, False, False), - ("demo_data/000000315905_zebra.jpg", 1e-6, 0, 0.7, 0.1, False, False), + ("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False), + ("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False), + ("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False), + ("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False), + ("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False), ), description=( "**Hugging Face Space (cpu-basic)** — lightweight demo: **CPU-only** PRIMA inference. " "The Space build skips Detectron2 and uses a full-image crop fallback. TTA is optional " - "(0 by default; increases runtime). Mesh `.obj` export is off by default to save time and disk." + "(30 iterations by default, matching the local demo; set to 0 to skip). Mesh `.obj` export " + "is off by default to save time and disk." ), interface_title="PRIMA on Hugging Face — lightweight CPU demo", ) @@ -806,6 +811,18 @@ def parse_args() -> argparse.Namespace: default=DEFAULT_OUT_FOLDER, help="Folder used to save rendered outputs and meshes", ) + parser.add_argument( + "--server_name", + type=str, + default=DEFAULT_SERVER_NAME, + help="Host/interface used by Gradio. Use 0.0.0.0 for Run:AI port-forward.", + ) + parser.add_argument( + "--server_port", + type=int, + default=DEFAULT_SERVER_PORT, + help="Port used by Gradio.", + ) return parser.parse_args() @@ -826,4 +843,9 @@ def parse_args() -> argparse.Namespace: out_folder=args.out_folder, runtime_cache=runtime_cache, ) - demo.launch(inbrowser=False, ssr_mode=False) + demo.launch( + inbrowser=False, + ssr_mode=False, + server_name=args.server_name, + server_port=args.server_port, + )