PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation

This is the official implementation of the approach described in the preprint:

PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
Xiaohang Yu, Ti Wang, Mackenzie Weygandt Mathis

---

TL;DR

PRIMA creates a 3D quadruped mesh from a single 2D image. It leverages BioCLIP-based biological priors for robust cross-species shape understanding, then applies test-time adaptation with 2D reprojection and auxiliary keypoint guidance to refine SMAL pose and shape predictions.

It further can be used to build Quadruped3D, a large-scale pseudo-3D dataset with diverse species and poses.

PRIMA achieves state-of-the-art results on Animal3D, CtrlAni3D, Quadruped2D, and Animal Kingdom datasets.

Installation

PRIMA requires Python 3.10 or newer. A CUDA-enabled PyTorch installation is recommended for local inference and training.

Install from PyPI

Create a clean environment, install PyTorch for your CUDA version, then install the package:

conda create -n prima python=3.10 -y
conda activate prima

# Example for CUDA 11.8. Adjust this command for your CUDA version.
pip install --index-url https://download.pytorch.org/whl/cu118 \
    "torch==2.2.1" "torchvision==0.17.1" "torchaudio==2.2.1"

# Install chumpy and PyTorch3D
python -m pip install --no-build-isolation \
      "git+https://github.com/mattloper/chumpy.git"
python -m pip install --no-build-isolation \
      "git+https://github.com/facebookresearch/pytorch3d.git"

# Install PRIMA from PyPI-test (for now)
pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple prima-animal==0.1.7

# Install PRIMA from PyPI
pip install prima-animal

prima-animal includes demo runtime dependencies used by demo.py, demo_tta.py, and app.py (including Detectron2 and DeepLabCut).

Install from this repository

Use this path if you want to run the code from a fresh clone.

git clone https://github.com/AdaptiveMotorControlLab/PRIMA.git
cd PRIMA

The helper script below creates a fresh virtual environment, installs runtime dependencies, pulls Git LFS assets if available, downloads the default demo checkpoints/data, and verifies that the demo dependencies can be imported:

PRIMA_PYTHON=/path/to/python3.10 \
PRIMA_VENV=prima_env \
./scripts/clean_install_local.sh
source prima_env/bin/activate

Options:

• --skip-data skips the large demo data download if data/ is already populated.
• --wipe-data --force-data removes downloaded demo assets and downloads them again.
• --no-editable installs dependencies without registering the repo as an editable package.

On macOS, install Python 3.10 if needed:

brew install python@3.10
PRIMA_PYTHON=/opt/homebrew/bin/python3.10 \
PRIMA_VENV=prima_env \
./scripts/clean_install_local.sh
source prima_env/bin/activate

If macOS reports Cannot read image: demo_data/..., install Git LFS and pull the demo images:

git lfs install
git lfs pull --include="demo_data/*"

---

Demo

Checkpoints and data

The demo scripts auto-download their default Stage 1 PRIMA assets from Hugging Face when the checkpoint or matching Hydra config is missing. If you want to pre-download all necessary checkpoints and data ahead of time, run:

python scripts/setup_demo_data.py --hf-repo-id MLAdaptiveIntelligence/PRIMA

Approximate default prefetch volume from Hugging Face is ~5.5 GB total (s1ckpt_inference.ckpt ~3 GB + amr_vitbb.pth ~2.5 GB + SMAL files). Expected time is roughly:

• 100 Mbps: ~7-10 minutes
• 300 Mbps: ~2-4 minutes
• 1 Gbps: ~1 minute

Existing files are reused by default; pass --force only if you need to redownload them. If you also need the Stage 3 pretrained model, add --include-stage3.

Expected files in that Hugging Face repo root:

• my_smpl_00781_4_all.pkl
• my_smpl_data_00781_4_all.pkl
• walking_toy_symmetric_pose_prior_with_cov_35parts.pkl
• amr_vitbb.pth
• config_s1_HYDRA.yaml
• s1ckpt_inference.ckpt

Optional Stage 3 prefetch expects:

• config_s3_HYDRA.yaml
• s3ckpt_inference.ckpt

Demo (without TTA)

Run animal detection + PRIMA 3D pose/shape inference:

bash demo.sh

Outputs are written to demo_out/. Edit demo.sh if you want to use a custom checkpoint path.

---

Demo (with TTA)

Run PRIMA inference with test-time adaptation:

bash demo_tta.sh

Outputs are written to demo_out_tta/ (before/after TTA renders, keypoints, and optional meshes). Edit demo_tta.sh if you want to change the checkpoint, TTA learning rate, or number of iterations.

---

Gradio demo

We also provide a simple Gradio-based web demo for interactive testing in the browser:

python app.py \
  --checkpoint data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt \
  --out_folder demo_out_tta_gradio/

This starts a local Gradio app (by default on http://127.0.0.1:7860), where you can upload images and visualize PRIMA predictions and adaptation results. The s1ckpt_inference.ckpt checkpoint is downloaded automatically if missing.

app.py picks a demo profile automatically:

	Local (python app.py)	Hugging Face Space
PRIMA device	GPU if available, else CPU	CPU only
Detector	Detectron2 X-101-FPN	DeepLabCut SuperAnimal detector
Default TTA iterations	30	30
Save .obj meshes	on	off
Preload checkpoint at startup	off	on

Override for testing: PRIMA_DEMO_MODE=local or PRIMA_DEMO_MODE=space.

---

Training and Evaluation

Dataset Setup

Download datasets from Animal3D, CtrlAni3D, Quadruped2D, and Animal Kingdom. For Quadruped2D, download the images from SuperAnimal-Quadruped80K and our processed annotations from here. Put all the datasets under datasets/.

Training

Two-stage training script:

bash train.sh

Training outputs are written to logs/train/runs/<exp_name>/.

Evaluation

python eval.py \
  --config data/PRIMAS1/.hydra/config.yaml \
  --checkpoint data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt

Common values for --dataset are controlled by:

• configs_hydra/experiment/default_val.yaml

---

Acknowledgements

This release builds on several open-source projects, including:

• Detectron2
• BioCLIP
• AniMer
• DeepLabCut
• SAM3DB

---

Citation

If you use this code in your research, please cite our PRIMA paper.

@misc{yu_prima,
  title={PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation},
  author={Xiaohang Yu and Ti Wang and Mackenzie Weygandt Mathis},
}

---

Contact

For issues, please open a GitHub issue in this repository.