BakedAvatar
Pytorch Code for "BakedAvatar: Baking Neural Fields for Real-Time Head Avatar Synthesis"
Install / Use
/learn @buaavrcg/BakedAvatarREADME
Paper | Project Page
<br> <tr> <img src="./assets/teaser.jpg" width="100%"/> </tr> </div>Setup
First, clone this repo:
git clone https://github.com/buaavrcg/BakedAvatar
cd BakedAvatar
Then, install the required environment. We recommend using Anaconda to manage your python environment. You can setup the required environment by the following commands:
conda env create -f environment.yml
conda activate BakedAvatar
Or you can setup the required environment manually:
conda create -n BakedAvatar python=3.10
conda activate BakedAvatar
# Install Pytorch (or follow specific instructions for your GPU on https://pytorch.org/get-started/locally/)
conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia
# Install various required libraries
pip install accelerate configargparse chumpy opencv-python pymeshlab trimesh scikit-image xatlas matplotlib tensorboard tqdm torchmetrics face-alignment
# Install Pytorch3D (Or follow instructions in https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md)
conda install pytorch3d -c pytorch3d # Linux only
# Install nvdiffrast
git clone https://github.com/NVlabs/nvdiffrast
cd nvdiffrast && pip install . && cd ..
# Install mise (for levelset extraction)
pip install Cython && pip install code/utils/libmise/
Finally, download FLAME model, choose FLAME 2020 and unzip it, copy 'generic_model.pkl' into ./code/flame/FLAME2020
Download Training Data
We use the same data format as in IMavatar and PointAvatar.
You can download a preprocessed dataset from subject 1, subject 2, then unzip the files into data/datasets folder. You should be able to see the paths of one subject's videos structured like data/datasets/<subject_name>/<video_name>.
To generate your own dataset, please follow the instructions in the IMavatar repo.
We also provide an example of the pre-trained checkpoint here.
Train implicit fields (Stage-1)
# Configure your training (use DDP?), see https://huggingface.co/docs/accelerate for details
accelerate config
cd code
accelerate launch scripts/runner.py -c config/subject1.yaml -t train
Bake meshes and textures (Stage-2)
# Extract the meshes
accelerate launch scripts/runner.py -c config/subject1.yaml -t mesh_export
# Precompute the textures and export MLP weights
# Note: change the path to the actual mesh_data.pkl path if you use a different config
accelerate launch scripts/runner.py -c config/subject1.yaml -t texture_export --mesh_data_path ../data/experiments/subject1/mesh_export/iter_30000/marching_cube/res_init16_up5/mesh_data.pkl
Fine tuning (Stage-3)
# Fine-tune the textures with higher resolution (512x512)
accelerate launch scripts/runner.py -c config/subject1.yaml -t fine_tuning --img_res 512 512 --batch_size 1 --mesh_data_path ../data/experiments/subject1/mesh_export/iter_30000/marching_cube/res_init16_up5/mesh_data.pkl
Run evaluation
# evaluate fine-tuned meshes (result of stage-3)
accelerate launch scripts/runner.py -c config/subject1.yaml -t test --img_res 512 512 --use_finetune_model --mesh_data_path ../data/experiments/subject1/finetune_mesh_data/iter_30000/mesh_data.pkl
# evaluate baked meshes (result of stage-2)
accelerate launch scripts/runner.py -c config/subject1.yaml -t test --img_res 512 512 --use_finetune_model --mesh_data_path ../data/experiments/subject1/mesh_export/iter_30000/marching_cube/res_init16_up5/mesh_data.pkl
# evaluate implicit fields (result of stage-1) (the rendering speed will be much slower)
accelerate launch scripts/runner.py -c config/subject1.yaml -t test --img_res 512 512
# run cross-identity reenactment using meshes in PyTorch code
# you may replace the reenact_data_dir with the path to the reenactment dataset and replace the reenact_subdirs with the subdirectories names
# if you would like to see reenactment results of implicit fields, remove --use_finetune_model
accelerate launch scripts/runner.py -c config/subject1.yaml -t test --img_res 512 512 \
--use_finetune_model --mesh_data_path ../data/experiments/subject1/finetune_mesh_data/iter_30000/mesh_data.pkl \
--reenact_data_dir ../data/datasets/soubhik --reenact_subdirs test
Export assets and run the real-time web demo
# export baked meshes and textures for the web demo
python scripts/unpack_pkl.py ../data/experiments/subject1/finetune_mesh_data/iter_30000/mesh_data.pkl --output ./mesh_data
# export the FLAME parameter sequence for reenactment
# The flame_params.json are from the files in the train and test subfolders of the dataset (e.g., ../data/datasets/soubhik/train/flame_params.json)
# You may export the sequences from the same identity for self-reenactment, or from different identities for cross-identity reenactment.
python scripts/export_flame_sequence.py <path to 1st flame_params.json> <path to 2nd flame_params.json> ... --output ./sequence_data
Copy the exported mesh_data directory and sequence_data directory into the root of the web demo and start the server.
The dictionary structure should be like:
web_demo
├── jsUtils
├── mesh_data
├── sequence_data
└── src
Make sure that you have installed Npm and Node.js, then run the following commands:
cd web_demo
npm install
npm run build
npm install --global serve
serve
Then, open your browser and visit http://localhost:8080/. To run the real-time reenactment, you can select one of the buttons with the name of the sequence in the web demo.
Citation
If you find our code or paper useful, please cite as:
@article{bakedavatar,
author = {Duan, Hao-Bin and Wang, Miao and Shi, Jin-Chuan and Chen, Xu-Chuan and Cao, Yan-Pei},
title = {BakedAvatar: Baking Neural Fields for Real-Time Head Avatar Synthesis},
year = {2023},
issue_date = {December 2023},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3618399},
doi = {10.1145/3618399},
volume = {42},
number = {6},
journal = {ACM Trans. Graph.},
month = {sep},
articleno = {225},
numpages = {14}
}
