Skip to content

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

License

Notifications You must be signed in to change notification settings

buaavrcg/BakedAvatar

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

8 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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

1State Key Laboratory of Virtual Reality Technology and Systems, Beihang University  2Zhongguancun Laboratory  3ARC Lab, Tencent PCG

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.

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}
}

About

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

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published