Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 40 additions & 50 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ Xiaohang Yu, Ti Wang, Mackenzie Weygandt Mathis
---


## 🚀 TL;DR
## 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.
Expand All @@ -21,15 +21,19 @@ PRIMA achieves state-of-the-art results on Animal3D, CtrlAni3D, Quadruped2D, and

## Installation

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

### Install from PyPI

> Recommended: Python 3.10 and a CUDA-enabled PyTorch installation.
Create a clean environment, install PyTorch for your CUDA version, then install
the package:

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

# Install PyTorch matching your CUDA (example: CUDA 11.8)
# 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"

Expand All @@ -39,48 +43,60 @@ python -m pip install --no-build-isolation \
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).

### Clean install from this repository
### Install from this repository

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

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

Use these when developing from a **git clone** (not the PyPI wheel). The shell scripts are **non-interactive** (pip uses `--no-input`; `GIT_TERMINAL_PROMPT=0` for git). Put Hugging Face credentials in your environment or git credential helper before pushing the Space.
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:

**Local (fresh venv, LFS assets, Hub demo weights, smoke test)** — requires **Python 3.10+**
(Gradio 5.1+ / Space-provided Gradio 6.x and `app.py` type hints). On macOS without `python3.10` on your `PATH`, install
`brew install python@3.10` and set `PRIMA_PYTHON=/opt/homebrew/bin/python3.10`.

```bash
chmod +x scripts/clean_install_local.sh scripts/clean_redeploy_hf_space.sh scripts/deploy_hf_space.sh
PRIMA_PYTHON=/opt/homebrew/bin/python3.10 ./scripts/clean_install_local.sh
PRIMA_PYTHON=/path/to/python3.10 \
PRIMA_VENV=prima_env \
./scripts/clean_install_local.sh
source prima_env/bin/activate
```

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.
- **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`.
- `--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.

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'`.
On macOS, install Python 3.10 if needed:

**Hugging Face Space (full redeploy from your working tree):**
```bash
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
```

Requires [Git LFS / Xet](https://huggingface.co/docs/hub/xet/using-xet-storage#git) tooling (`brew install git-lfs git-xet`, `git xet install`, `git lfs install`). Then:
If macOS reports `Cannot read image: demo_data/...`, install Git LFS and pull
the demo images:

```bash
./scripts/clean_redeploy_hf_space.sh
git lfs install
git lfs pull --include="demo_data/*"
```

This is equivalent to `./scripts/deploy_hf_space.sh` and force-pushes a fresh snapshot to the Space.

---

## Demo
Expand Down Expand Up @@ -170,32 +186,6 @@ The `s1ckpt_inference.ckpt` checkpoint is downloaded automatically if missing.

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

#### Hugging Face Space (maintainers)

Demo images under `demo_data/` and `images/teaser.png` are tracked with **Git LFS**
(see `.gitattributes`) so they can be pushed to a Hugging Face Space under the Hub’s
LFS / **Xet** bridge. Install tooling once:

```bash
brew install git-lfs git-xet
git xet install
git lfs install
```

Then from a clean checkout with LFS files present, redeploy the Space (same as `clean_redeploy_hf_space.sh`):

```bash
./scripts/deploy_hf_space.sh
# or
./scripts/clean_redeploy_hf_space.sh
```

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 the DeepLabCut SuperAnimal detector fallback on the CPU Space.

---


Expand Down
70 changes: 60 additions & 10 deletions scripts/clean_install_local.sh
Original file line number Diff line number Diff line change
Expand Up @@ -68,29 +68,69 @@ resolve_python() {
return 1
fi
local c p
for c in python3.12 python3.11 python3.10; do
for c in python3.10 python3.11; do
if command -v "$c" >/dev/null 2>&1; then
if "$c" -c 'import sys; raise SystemExit(0 if sys.version_info >= (3, 10) else 1)'; then
command -v "$c"
return 0
fi
fi
done
for p in /opt/homebrew/bin/python3.10 /usr/local/bin/python3.10; do
for p in /opt/homebrew/bin/python3.10 /usr/local/bin/python3.10 /opt/homebrew/opt/python@3.10/bin/python3.10 /usr/local/opt/python@3.10/bin/python3.10; do
if [[ -x "$p" ]]; then
echo "$p"
return 0
fi
done
if command -v python3.12 >/dev/null 2>&1; then
command -v python3.12
return 0
fi
return 1
}

resolve_torch_index_url() {
if [[ -n "${PRIMA_TORCH_INDEX_URL:-}" ]]; then
echo "${PRIMA_TORCH_INDEX_URL}"
return 0
fi

if [[ "$(uname -s)" == "Darwin" ]]; then
return 1
fi

if command -v nvcc >/dev/null 2>&1; then
local cuda_version
cuda_version="$(nvcc --version | sed -n 's/.*release \([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' | head -n 1)"
case "${cuda_version}" in
11.8)
echo "https://download.pytorch.org/whl/cu118"
return 0
;;
12.1)
echo "https://download.pytorch.org/whl/cu121"
return 0
;;
"")
echo "[clean-install] WARN: Could not parse nvcc CUDA version; using pip default PyTorch wheel." >&2
return 1
;;
*)
echo "[clean-install] WARN: CUDA ${cuda_version} detected; set PRIMA_TORCH_INDEX_URL if Detectron2 needs a specific PyTorch wheel." >&2
return 1
;;
esac
fi

return 1
}

echo "[clean-install] Repository: ${ROOT}"

if ! PY="$(resolve_python)"; then
echo "[clean-install] ERROR: Need Python 3.10 or newer (Gradio 5 + app type hints)." >&2
echo " macOS: brew install python@3.10" >&2
echo " Then: PRIMA_PYTHON=/opt/homebrew/bin/python3.10 $0 ..." >&2
echo "[clean-install] ERROR: Need Python 3.10, 3.11, or 3.12 on PATH, or set PRIMA_PYTHON." >&2
echo " Conda example: PRIMA_PYTHON=\$HOME/miniconda3/envs/prima/bin/python3.10 $0 ..." >&2
echo " macOS example: brew install python@3.10 && PRIMA_PYTHON=/opt/homebrew/bin/python3.10 $0 ..." >&2
exit 1
fi
echo "[clean-install] Using Python: $("$PY" -c 'import sys; print(sys.executable, sys.version.split()[0])')"
Expand All @@ -116,10 +156,20 @@ source "${VENV}/bin/activate"
python -m pip install --no-input -U pip wheel
# Match requirements.txt / pyproject pins before pulling the rest
python -m pip install --no-input "setuptools<81" "packaging<25" "Cython<3"
python -m pip install --no-input "numpy==1.26.1"

echo "[clean-install] xtcocotools (needs numpy available during build) ..."
python -m pip install --no-input --no-build-isolation "xtcocotools==1.14.3"

if TORCH_INDEX_URL="$(resolve_torch_index_url)"; then
echo "[clean-install] Installing PyTorch from ${TORCH_INDEX_URL} ..."
python -m pip install --no-input --index-url "${TORCH_INDEX_URL}" \
"torch==2.2.1" "torchvision==0.17.1"
fi

echo "[clean-install] pip install -r requirements.txt (this can take a long time) ..."
REQ_TMP="$(mktemp)"
grep -vE '^[[:space:]]*(deeplabcut|detectron2)' "${ROOT}/requirements.txt" > "${REQ_TMP}"
grep -vE '^[[:space:]]*(deeplabcut|detectron2|xtcocotools)' "${ROOT}/requirements.txt" > "${REQ_TMP}"
python -m pip install --no-input -r "${REQ_TMP}"
rm -f "${REQ_TMP}"

Expand Down Expand Up @@ -153,12 +203,12 @@ if [[ "$WIPE_DATA" -eq 1 ]]; then
fi

if [[ "$SKIP_DATA" -eq 0 ]]; then
FORCE_ARGS=()
echo "[clean-install] Downloading demo assets (large) ..."
if [[ "$FORCE_DATA" -eq 1 ]]; then
FORCE_ARGS=(--force)
python "${ROOT}/scripts/setup_demo_data.py" --force
else
python "${ROOT}/scripts/setup_demo_data.py"
fi
echo "[clean-install] Downloading demo assets (large) ..."
python "${ROOT}/scripts/setup_demo_data.py" "${FORCE_ARGS[@]}"
else
echo "[clean-install] Skipping setup_demo_data (--skip-data)."
fi
Expand Down
Loading