Code release for our preprint:
Brent Yi1, Vickie Ye1, Maya Zheng1, Lea Müller1, Georgios Pavlakos2, Yi Ma1, Jitendra Malik1, and Angjoo Kanazawa1. Estimating Body and Hand Motion in an Ego-sensed World. arXiV, 2024. |
1UC Berkeley, 2UT Austin
TLDR; We use egocentric SLAM poses and images to estimate 3D human body pose, height, and hands.
https://github.com/user-attachments/assets/7d28e07f-ab83-4749-ac6b-abe692d9ba20
This repository is structured as follows:
.
├── download_checkpoint_and_data.sh
│ - Download model checkpoint and sample data.
├── 0_preprocess_training_data.py
│ - Preprocessing script for training datasets.
├── 1_train_motion_prior.py
│ - Training script for motion diffusion model.
├── 2_run_hamer_on_vrs.py
│ - Run HaMeR on inference data (expects Aria VRS).
├── 3_aria_inference.py
│ - Run full pipeline on inference data.
├── 4_visualize_outputs.py
│ - Visualize outputs from inference.
│
├── src/egoallo/
│ ├── data/ - Dataset utilities.
│ ├── transforms/ - SO(3) / SE(3) transformation helpers.
│ └── *.py - All core implementation.
│
└── pyproject.toml - Python dependencies/package metadata.
EgoAllo requires Python 3.12 or newer.
Clone the repository.
git clone https://github.com/brentyi/egoallo.git
Install general dependencies.
cd egoallo
pip install -e .
Download+unzip model checkpoint and sample data.
bash download_checkpoint_and_data.sh
You can also download the zip files manually: here are links to the checkpoint and example trajectories.
Download the SMPL-H model file.
You can find the "Extended SMPL+H model" from the MANO project webpage.
Our scripts assumes an npz file located at ./data/smplh/neutral/model.npz
, but this can be overridden at the command-line (--smplh-npz-path {your path}
).
Visualize model outputs.
The example trajectories directory includes example outputs from our model. You can visualize them with:
python 4_visualize_outputs.py --search-root-dir ./egoallo_example_trajectories
Installing inference dependencies.
Our guidance optimization uses a Levenberg-Marquardt optimizer that's implemented in JAX. If you want to run this on an NVIDIA GPU, you'll need to install JAX with CUDA support:
# Also see: https://jax.readthedocs.io/en/latest/installation.html
pip install -U "jax[cuda12]"
You'll also need jaxls:
pip install git+https://github.com/brentyi/jaxls.git
Running inference on example data.
Here's an example command for running EgoAllo on the "coffeemachine" sequence:
python 3_aria_inference.py --traj-root ./egoallo_example_trajectories/coffeemachine
You can run python 3_aria_inference.py --help
to see the full list of options.
Running inference on your own data.
To run inference on your own data, you can copy the structure of the example trajectories. The key files are:
closed_loop_trajectory.csv
and semidense_points.csv.gz
.hamer_outputs.pkl
.Running HaMeR on your own data.
To generate the hamer_outputs.pkl
file, you'll need to install hamer_helper.
Then, as an example for running on our coffeemachine sequence:
python 2_run_hamer_on_vrs.py --traj-root ./egoallo_example_trajectories/coffeemachine
This repository currently contains:
egoallo
package, which contains reference training and sampling implementation details.While we've put effort into cleaning up our code for release, this is research code and there's room for improvement. If you have questions or comments, please reach out!