BakedAvatar: Baking Neural Fields for Real-Time Head Avatar Synthesis

Hao-Bin Duan¹ Miao Wang^{1, 2} Jin-Chuan Shi¹ Xu-Chuan Chen¹ Yan-Pei Cao³

¹State Key Laboratory of Virtual Reality Technology and Systems, Beihang University ²Zhongguancun Laboratory ³ARC Lab, Tencent PCG

ACM Transactions on Graphics (SIGGRAPH Asia 2023)

TL;DR

BakedAvatar takes monocular video recordings of a person and produces a mesh-based representation for real-time 4D head avatar synthesis on various devices including mobiles.

### [Paper](https://dl.acm.org/doi/10.1145/3618399) | [Project Page](https://buaavrcg.github.io/BakedAvatar/)

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)
accelerate launch scripts/runner.py -c config/subject1.yaml -t test --img_res 512 512

# run cross-identity reenactment 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
accelerate launch scripts/runner.py -c config/subject1.yaml -t test --img_res 512 512 --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}
}

buaavrcg / BakedAvatar

readme