""" PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation Official implementation of the paper: "PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation" by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis Licensed under a modified MIT license """ """Gradio demo for PRIMA + SuperAnimal + TTA. This script wraps the ``demo_tta.py`` pipeline into an interactive Gradio interface. The overall logic follows: 1. Given an input image, run Detectron2 to detect animals. 2. For each detected animal, run PRIMA for 3D pose/shape estimation. 3. Run the fine-tuned DeepLabCut SuperAnimal model to obtain PRIMA 26-keypoint 2D predictions. 4. Run test-time adaptation (TTA) with user-specified lr and iters. 5. Render and save before/after TTA results and keypoint visualizations. """ 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 Any, Callable, Dict, List, Optional, Tuple from pathlib import Path # macOS: keep compute single-threaded and run inference on main thread. if sys.platform == "darwin": os.environ.setdefault("OMP_NUM_THREADS", "1") import cv2 import gradio as gr import numpy as np import torch import torch.utils.data if sys.platform == "darwin": torch.set_num_threads(1) # Repo-local minimal ``chumpy`` shim (see ``chumpy/__init__.py``) so SMAL pickles load # without installing the full chumpy package in Space builds. _REPO_ROOT = Path(__file__).resolve().parent if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) from prima.utils.weights import ( DEFAULT_HF_REPO_ID, resolve_prima_checkpoint_path, ) from prima.utils.detection import select_animal_boxes # Default checkpoint path following README instructions DEFAULT_CHECKPOINT = str(_REPO_ROOT / "data" / "PRIMAS1" / "checkpoints" / "s1ckpt_inference.ckpt") DEFAULT_HF_ASSET_REPO = DEFAULT_HF_REPO_ID # 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 = ( "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/" "faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl" ) _D2_X101_CFG = "COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml" _D2_X101_URL = ( "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/" "faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl" ) # Gradio example row: (image_rel, tta_lr, tta_iters, det_thresh, kp_thresh, side_view, render_depth, save_mesh) ExampleRow = Tuple[str, float, int, float, float, bool, bool, bool] @dataclass(frozen=True) class DemoProfile: """Runtime settings for either the full local app or the lightweight HF Space demo.""" mode: str prima_device: str # "auto" (CUDA if available) or "cpu" detectron_config_yaml: str detectron_weights_url: str detectron_device: str # "auto" or "cpu" default_tta_iters: int max_tta_iters: int default_save_mesh: bool default_side_view: bool default_render_depth: bool preload_assets: bool example_rows: Tuple[ExampleRow, ...] description: str interface_title: str def resolve_prima_device(self) -> torch.device: if self.prima_device == "cpu": return torch.device("cpu") return torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") def resolve_detectron_device(self) -> str: if self.detectron_device == "cpu": return "cpu" return "cuda" if torch.cuda.is_available() else "cpu" LOCAL_DEMO_PROFILE = DemoProfile( mode="local", prima_device="auto", detectron_config_yaml=_D2_X101_CFG, detectron_weights_url=_D2_X101_URL, detectron_device="auto", default_tta_iters=30, max_tta_iters=100, default_save_mesh=True, default_side_view=False, default_render_depth=False, preload_assets=False, example_rows=( ("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False, True), ("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False, True), ("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False, True), ("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False, True), ("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False, True), ), description=( "**Local demo** — full pipeline on your machine (GPU when available).\n\n" "Detectron2 **X-101-FPN**, PRIMA mesh recovery, optional **DeepLabCut SuperAnimal + TTA**. " "Set TTA iterations to **0** to skip adaptation. Outputs are saved under " f"`{DEFAULT_OUT_FOLDER}`." ), interface_title=( "PRIMA local demo (GPU/CPU) — detection, mesh recovery, optional TTA" ), ) SPACE_DEMO_PROFILE = DemoProfile( mode="space", prima_device="cpu", detectron_config_yaml=_D2_R50_CFG, detectron_weights_url=_D2_R50_URL, detectron_device="cpu", default_tta_iters=30, max_tta_iters=30, default_save_mesh=False, default_side_view=False, default_render_depth=False, preload_assets=True, example_rows=( ("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False, False), ("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False, False), ("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False, False), ("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False, False), ("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False, False), ), description=( "**Hugging Face Space (cpu-basic)** — lightweight demo: **CPU-only** PRIMA inference. " "The Space build skips Detectron2 and uses the DeepLabCut SuperAnimal detector for animal " "crops. TTA is optional (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", ) def _is_truthy_env(var_name: str) -> bool: return os.environ.get(var_name, "").strip().lower() in {"1", "true", "yes", "on"} def _running_on_space() -> bool: return bool(os.environ.get("SPACE_ID") or os.environ.get("HF_SPACE_ID")) @lru_cache(maxsize=1) def get_demo_profile() -> DemoProfile: """Select local vs Space profile. Override with ``PRIMA_DEMO_MODE=local|space``.""" override = os.environ.get("PRIMA_DEMO_MODE", "").strip().lower() if override == "local": return LOCAL_DEMO_PROFILE if override == "space": return SPACE_DEMO_PROFILE return SPACE_DEMO_PROFILE if _running_on_space() else LOCAL_DEMO_PROFILE def _gradio_examples_for_interface(profile: DemoProfile) -> List[List]: """Gradio prefetches example media at startup (paths must exist beside ``app.py``).""" if _is_truthy_env("PRIMA_DISABLE_GRADIO_EXAMPLES"): return [] rows: List[List] = [] for rel, *rest in profile.example_rows: p = _REPO_ROOT / rel if p.is_file(): rows.append([str(p), *rest]) return rows def _should_use_gradio_queue(profile: DemoProfile) -> bool: """Whether to enable Gradio's background queue. On macOS local, the queue runs handlers in worker threads; PyRender and pyglet/AppKit then crash with "setting the main menu on a non-main thread". """ if _is_truthy_env("PRIMA_GRADIO_NO_QUEUE"): return False if _is_truthy_env("PRIMA_GRADIO_QUEUE"): return True return not (sys.platform == "darwin" and profile.mode == "local") def _warmup_runtime_cache(checkpoint_path: str, profile: DemoProfile) -> Dict[str, Any]: """Load model + DLC on the main thread (recommended for macOS local).""" print("[startup] Preloading DeepLabCut on main thread…") _deeplabcut_available() print("[startup] Loading PRIMA + Detectron2 on main thread (first run can take several minutes)…") model, model_cfg, renderer, cam_crop_to_full_fn, device, detector = _load_model_and_detector_for_demo( checkpoint_path, profile ) return { "model": model, "model_cfg": model_cfg, "renderer": renderer, "cam_crop_to_full_fn": cam_crop_to_full_fn, "device": device, "detector": detector, } def _should_preload_assets(profile: DemoProfile) -> bool: preload_env = os.environ.get("PRIMA_PRELOAD_ASSETS") if preload_env is not None: return _is_truthy_env("PRIMA_PRELOAD_ASSETS") return profile.preload_assets def _deeplabcut_available() -> bool: try: from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images # noqa: F401 return True except Exception: return False def _preload_assets_once(checkpoint_path: str) -> None: print("[startup] Ensuring demo assets from Hugging Face Hub...") resolve_prima_checkpoint_path( checkpoint_path, data_dir=_REPO_ROOT / "data", auto_download=True, hf_repo_id=os.environ.get("PRIMA_HF_REPO_ID", DEFAULT_HF_ASSET_REPO), ) print("[startup] Asset preload complete.") def _load_prima_model(checkpoint_path: str = DEFAULT_CHECKPOINT): """Load PRIMA model and renderer once for the Gradio app.""" from prima.models import load_prima from prima.utils.renderer import Renderer, cam_crop_to_full checkpoint_path = resolve_prima_checkpoint_path( checkpoint_path, data_dir=_REPO_ROOT / "data", auto_download=True, hf_repo_id=os.environ.get("PRIMA_HF_REPO_ID", DEFAULT_HF_ASSET_REPO), ) checkpoint = Path(checkpoint_path) cfg_path = checkpoint.parent.parent / ".hydra" / "config.yaml" if not checkpoint.exists(): raise FileNotFoundError( f"Missing checkpoint: {checkpoint}. Download demo checkpoints/data as described in README." ) if not cfg_path.exists(): raise FileNotFoundError( f"Missing model config: {cfg_path}. Ensure the full checkpoint folder layout from README is present." ) profile = get_demo_profile() model, model_cfg = load_prima(checkpoint_path) device = profile.resolve_prima_device() model = model.to(device) model.eval() renderer = Renderer(model_cfg, faces=model.smal.faces) return model, model_cfg, renderer, cam_crop_to_full, device def _build_detector(profile: Optional[DemoProfile] = None): """Build Detectron2 animal detector (profile selects X-101+GPU locally vs R50+CPU on Space).""" try: import detectron2.config import detectron2.engine from detectron2 import model_zoo except Exception as e: print(f"[warn] Detectron2 unavailable ({type(e).__name__}: {e}); using SuperAnimal detector fallback.") return None if profile is None: profile = get_demo_profile() config_yaml = profile.detectron_config_yaml weights = profile.detectron_weights_url device_str = profile.resolve_detectron_device() print(f"[detectron2] mode={profile.mode} config={config_yaml} device={device_str}") cfg = detectron2.config.get_cfg() cfg.merge_from_file(model_zoo.get_config_file(config_yaml)) cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 cfg.MODEL.WEIGHTS = weights cfg.MODEL.DEVICE = device_str detector = detectron2.engine.DefaultPredictor(cfg) return detector def _filter_superanimal_boxes( payload: Dict[str, Any], det_thresh: float, img_shape: Tuple[int, int], ) -> Optional[np.ndarray]: boxes = payload.get("bboxes") if boxes is None: return None boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4) if len(boxes) == 0: return None scores = payload.get("bbox_scores") if scores is None: scores = np.ones((len(boxes),), dtype=np.float32) scores = np.asarray(scores, dtype=np.float32).reshape(-1) if len(scores) != len(boxes): scores = np.ones((len(boxes),), dtype=np.float32) h, w = img_shape valid = ( (scores > float(det_thresh)) & np.isfinite(boxes).all(axis=1) & (boxes[:, 2] > 0.0) & (boxes[:, 3] > 0.0) ) if not np.any(valid): return None xywh = boxes[valid] scores = scores[valid] boxes = xywh.copy() boxes[:, 2] = xywh[:, 0] + xywh[:, 2] boxes[:, 3] = xywh[:, 1] + xywh[:, 3] boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0.0, float(max(1, w - 1))) boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0.0, float(max(1, h - 1))) valid_size = (boxes[:, 2] > boxes[:, 0]) & (boxes[:, 3] > boxes[:, 1]) boxes = boxes[valid_size] scores = scores[valid_size] if len(boxes) == 0: return None order = np.argsort(scores)[::-1] return boxes[order].astype(np.float32, copy=False) def _detect_superanimal_boxes(img_rgb: np.ndarray, det_thresh: float) -> Optional[np.ndarray]: try: from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images except Exception as e: print(f"[warn] DeepLabCut SuperAnimal unavailable ({type(e).__name__}: {e}); no fallback bbox.") return None with tempfile.TemporaryDirectory(prefix="sa_detect_") as tmp_dir: img_path = os.path.join(tmp_dir, "image.png") cv2.imwrite(img_path, cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)) dlc_device = "cuda" if torch.cuda.is_available() else "cpu" preds = superanimal_analyze_images( superanimal_name=SUPER_ANIMAL_ARGS.superanimal_name, model_name=SUPER_ANIMAL_ARGS.superanimal_model_name, detector_name=SUPER_ANIMAL_ARGS.superanimal_detector_name, images=img_path, max_individuals=SUPER_ANIMAL_ARGS.superanimal_max_individuals, out_folder=tmp_dir, progress_bar=False, device=dlc_device, pose_threshold=0.1, bbox_threshold=float(det_thresh), plot_skeleton=False, ) payload = preds.get(img_path) if payload is None: return None return _filter_superanimal_boxes(payload, det_thresh, img_rgb.shape[:2]) def _load_model_and_detector_for_demo(checkpoint_path: str, profile: DemoProfile): """Load PRIMA and Detectron2 once for the Gradio session (main thread only).""" model, model_cfg, renderer, cam_crop_to_full_fn, device = _load_prima_model(checkpoint_path) detector = _build_detector(profile) return model, model_cfg, renderer, cam_crop_to_full_fn, device, detector def _detect_animal_boxes( detector, img_bgr: np.ndarray, det_thresh: float, ) -> Optional[np.ndarray]: """Return Nx4 XYXY boxes or None if no animal detections.""" if detector is None: img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) return _detect_superanimal_boxes(img_rgb, det_thresh) 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 return boxes # SuperAnimal defaults (same as in demo_tta parser) SUPER_ANIMAL_ARGS = SimpleNamespace( superanimal_name="superanimal_quadruped", superanimal_model_name="hrnet_w32", superanimal_detector_name="fasterrcnn_resnet50_fpn_v2", superanimal_max_individuals=1, saved_2d_model_path="", pytorch_config_2d_path=str(_REPO_ROOT / "configs" / "sa_finetune_hrnet_w32.yaml"), ) def _collect_animal_results( model, model_cfg, renderer, cam_crop_to_full_fn, device, detector, out_folder: str, img_rgb: np.ndarray, tta_lr: float, tta_num_iters: int, det_thresh: float, kp_conf_thresh: float, side_view: bool, render_depth: bool, save_mesh: bool, 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. Returns: before_imgs: list of HxWx3 RGB images (before TTA) for all animals after_imgs: list of HxWx3 RGB images (after TTA) for all animals kpt_imgs: list of HxWx3 RGB keypoint visualizations first_before_mesh: path to first animal's before-TTA mesh (.obj) or None first_after_mesh: path to first animal's after-TTA mesh (.obj) or None """ from prima.utils import recursive_to from prima.datasets.vitdet_dataset import ViTDetDataset from demo_tta import ( denorm_patch_to_rgb, resolve_sa_weights_path, run_superanimal_on_patch, save_keypoint_vis, 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("") img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) if boxes is None: if detector is None: report("Detectron2 unavailable; detecting animals with SuperAnimal...") 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) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0) before_imgs: List[np.ndarray] = [] after_imgs: List[np.ndarray] = [] kpt_imgs: List[np.ndarray] = [] before_mesh_paths: List[str] = [] after_mesh_paths: List[str] = [] img_token = next(tempfile._get_candidate_names()) 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) animal_id = int(batch["animalid"][0]) # Save/render before TTA 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, out_before, batch, img_fn, animal_id, out_folder, suffix="before_tta", side_view=side_view, save_mesh=save_mesh, render_depth=render_depth, ) before_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_before_tta.png") if os.path.exists(before_png_path): before_bgr = cv2.imread(before_png_path) if before_bgr is not None: before_imgs.append(cv2.cvtColor(before_bgr, cv2.COLOR_BGR2RGB)) if save_mesh: before_obj_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_before_tta.obj") if os.path.exists(before_obj_path): 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, out_before, batch, img_fn, animal_id, out_folder, suffix="after_tta", side_view=side_view, save_mesh=save_mesh, render_depth=render_depth, ) after_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.png") if os.path.exists(after_png_path): after_bgr = cv2.imread(after_png_path) if after_bgr is not None: after_imgs.append(cv2.cvtColor(after_bgr, cv2.COLOR_BGR2RGB)) if save_mesh: after_obj_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.obj") if os.path.exists(after_obj_path): after_mesh_paths.append(after_obj_path) 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) if bodyparts_xyc is None: # No keypoints => skip TTA for this animal continue kpts_xyc = bodyparts_xyc kpts_xyc[kpts_xyc[:, 2] < float(kp_conf_thresh), 2] = 0.0 # Save keypoint visualization and npy kpt_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_prima26_kpts.png") save_keypoint_vis(patch_rgb, kpts_xyc, kpt_png_path) npy_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_prima26_kpts.npy") np.save(npy_path, kpts_xyc) if os.path.exists(kpt_png_path): kpt_bgr = cv2.imread(kpt_png_path) if kpt_bgr is not None: kpt_imgs.append(cv2.cvtColor(kpt_bgr, cv2.COLOR_BGR2RGB)) # Normalize keypoints to [-0.5, 0.5] as in demo_tta patch_h, patch_w = patch_rgb.shape[:2] kpts_norm = kpts_xyc.copy() kpts_norm[:, 0] = kpts_norm[:, 0] / float(patch_w) - 0.5 kpts_norm[:, 1] = kpts_norm[:, 1] / float(patch_h) - 0.5 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, gt_kpts_norm, num_iters=int(tta_num_iters), lr=float(tta_lr), ) report(f"Animal {batch_idx}/{total_batches}: rendering after TTA...") render_and_save( renderer, cam_crop_to_full_fn, out_after, batch, img_fn, animal_id, out_folder, suffix="after_tta", side_view=side_view, save_mesh=save_mesh, render_depth=render_depth, ) after_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.png") if os.path.exists(after_png_path): after_bgr = cv2.imread(after_png_path) if after_bgr is not None: after_imgs.append(cv2.cvtColor(after_bgr, cv2.COLOR_BGR2RGB)) if save_mesh: after_obj_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.obj") if os.path.exists(after_obj_path): after_mesh_paths.append(after_obj_path) 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 def build_demo( checkpoint_path: str = DEFAULT_CHECKPOINT, out_folder: str = DEFAULT_OUT_FOLDER, runtime_cache: Optional[Dict[str, Any]] = None, ) -> gr.Interface: profile = get_demo_profile() print( f"[demo] profile={profile.mode} prima={profile.resolve_prima_device()} " f"detectron={profile.detectron_config_yaml} d2_device={profile.resolve_detectron_device()}" ) if _should_use_gradio_queue(profile): print("[demo] Gradio queue enabled (background worker threads).") else: print("[demo] Gradio queue disabled (inference runs on main thread; required on macOS local).") os.makedirs(out_folder, exist_ok=True) runtime_cache = runtime_cache or { "model": None, "model_cfg": None, "renderer": None, "cam_crop_to_full_fn": None, "device": None, "detector": None, } def gradio_inference( image: np.ndarray, tta_lr: float, tta_num_iters: int, det_thresh: float, kp_conf_thresh: float, side_view: bool, render_depth: bool, save_mesh: bool, ): """Wrapper for Gradio. ``image`` is an RGB numpy array. Yields intermediate status so long first-run (Hub downloads + model load) and long inference do not hit silent client/proxy WebSocket timeouts. """ if image is None: yield None, None, None, "No image provided." return if int(tta_num_iters) > 0 and not _deeplabcut_available(): yield ( None, None, None, "DeepLabCut is not installed. Set **TTA iterations** to **0** for PRIMA-only inference, " "or install `deeplabcut` (see README / requirements.txt).", ) return if image.dtype != np.uint8: img_rgb = np.clip(image, 0, 255).astype(np.uint8) else: img_rgb = image yield None, None, None, "Queued; preparing run…" if runtime_cache["model"] is None: yield ( None, None, None, "First run: downloading demo assets from Hugging Face (large checkpoint) " "and loading the model. This can take many minutes.", ) try: model, model_cfg, renderer, cam_crop_to_full_fn, device, detector = _load_model_and_detector_for_demo( checkpoint_path, profile ) except Exception: yield None, None, None, f"Model initialization failed:\n{traceback.format_exc()}" return runtime_cache["model"] = model runtime_cache["model_cfg"] = model_cfg runtime_cache["renderer"] = renderer runtime_cache["cam_crop_to_full_fn"] = cam_crop_to_full_fn runtime_cache["device"] = device runtime_cache["detector"] = detector yield None, None, None, "Model loaded." try: yield None, None, None, "Running animal detection…" img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) boxes = _detect_animal_boxes(runtime_cache["detector"], img_bgr, det_thresh) if boxes is None: yield ( None, None, None, "No animal detected. Try lowering the detection threshold or another image.", ) 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, render_depth=render_depth, 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: stage_updates.put(message) with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: fut = pool.submit( run_collect, 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=1.0 ) break except concurrent.futures.TimeoutError: elapsed = int(time.monotonic() - t0) yield None, None, None, ( f"{latest_stage}\n" 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 first_before = before_imgs[0] if before_imgs else None first_after = after_imgs[0] if after_imgs else None first_kpts = kpt_imgs[0] if kpt_imgs else None if first_before is None and first_after is None: yield ( None, None, None, "No output generated. Try an image with a clearly visible quadruped.", ) return yield first_before, first_after, first_kpts, "OK" _gradio_examples = _gradio_examples_for_interface(profile) _iface_kw = dict( fn=gradio_inference, analytics_enabled=False, cache_examples=False, inputs=[ gr.Image( label="Input image", type="numpy", sources=["upload", "clipboard"], ), gr.Slider( label="TTA learning rate", minimum=1e-7, maximum=1e-4, value=1e-6, step=1e-7, ), gr.Slider( label="TTA iterations", minimum=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.", ), gr.Slider( label="Detection threshold", minimum=0.3, maximum=0.9, value=0.7, step=0.05, ), gr.Slider( label="Keypoint confidence threshold", minimum=0.0, maximum=1.0, value=0.1, step=0.05, ), gr.Checkbox(label="Render side view", value=profile.default_side_view), gr.Checkbox(label="Render depth map", value=profile.default_render_depth), gr.Checkbox(label="Save meshes (.obj)", value=profile.default_save_mesh), ], outputs=[ gr.Image(label="Before TTA"), gr.Image(label="After TTA"), gr.Image(label="PRIMA 26 keypoints"), gr.Textbox(label="Status / Traceback", lines=12), ], title=profile.interface_title, description=profile.description, ) if _gradio_examples: _iface_kw["examples"] = _gradio_examples demo = gr.Interface(**_iface_kw) if _should_use_gradio_queue(profile): demo.queue(max_size=8, default_concurrency_limit=1) return demo def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Gradio demo for PRIMA + SuperAnimal + TTA") parser.add_argument( "--checkpoint", type=str, default=DEFAULT_CHECKPOINT, help="Path to the pretrained PRIMA checkpoint", ) parser.add_argument( "--out_folder", type=str, 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() if __name__ == "__main__": args = parse_args() profile = get_demo_profile() if _should_preload_assets(profile): _preload_assets_once(args.checkpoint) 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, server_name=args.server_name, server_port=args.server_port, )