Home

Awesome

HeadNeRF: A Real-time NeRF-based Parametric Head Model

<span id="code_demo"></span>

This repository contains a pytorch implementation of "HeadNeRF: A Real-time NeRF-based Parametric Head Model (CVPR 2022)". Authors: Yang Hong, Bo Peng, Haiyao Xiao, Ligang Liu and Juyong Zhang*.

| Project Page | Paper |

This code has been tested on ubuntu 20.04/18.04 and contains the following parts:

  1. An interactive GUI that allows users to utilize HeadNeRF to directly edit the generated images’ rendering pose and various semantic attributes.
  2. A fitting framework for obtaining the latent code embedding in HeadNeRF of a single image.

Requirements

<!-- conda install -c bottler nvidiacub conda install -c fvcore -c iopath -c conda-forge fvcore iopath -->

Note:

<!-- - For NVIDIA 30 series GPUs, installing the pytorch with the specified version may be necessary to avoid compatibility issues. Please refer to [pytorch](https://pytorch.org/get-started/locally/) for details. --> <!-- ``` pip3 install torch==1.8.2+cu111 torchvision==0.9.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html ``` --> <!-- By testing, it is found that the following command is working for installing pytorch on NVIDIA 30 series GPUs. -->

Getting Started

Download ConfigModels.zip, TrainedModels.zip, and LatentCodeSamples.zip, then unzip them to the root dir of this project.

Other links: Google Drive, One Drive

The folder structure is as follows:

headnerf
├── ConfigModels
│   ├── faceparsing_model.pth
│   ├── nl3dmm_dict.pkl
│   └── nl3dmm_net_dict.pth
│
├── TrainedModels
│   ├── model_Reso32.pth
│   ├── model_Reso32HR.pth
│   └── model_Reso64.pth
│
└── LatentCodeSamples
    ├── model_Reso32
    │   ├── S001_E01_I01_P02.pth
    │   └── ...
    ├── model_Reso32HR
    │   ├── S001_E01_I01_P02.pth
    │   └── ...
    └── model_Reso64
        ├── S001_E01_I01_P02.pth
        └── ...

Note:

The Interactive GUI

#GUI, for editing the generated images’ rendering pose and various semantic attributes.
python MainGUI.py --model_path "TrainedModels/model_Reso64.pth"

Args:

An interactive interface like the first figure of this document will be generated after executing the above command.

<!-- ![](docs/Demo-Gui.png) --> <!-- <img src='docs/main-gui.png' align='center' style=' BORDER:#000000 1px solid'/> -->

The fitting framework

This part provides a framework for fitting a single image using HeadNeRF. Besides, some test images are provided in test_data/single_images dir. These images are from FFHQ dataset and do not participate in building HeadNeRF's models.

Data Preprocess

# generating head's mask.
python DataProcess/Gen_HeadMask.py --img_dir "test_data/single_images"

# generating 68-facial-landmarks by face-alignment, which is from 
# https://github.com/1adrianb/face-alignment
python DataProcess/Gen_Landmark.py --img_dir "test_data/single_images"

# generating the 3DMM parameters
python Fitting3DMM/FittingNL3DMM.py --img_size 512 \
                                    --intermediate_size 256  \
                                    --batch_size 9 \
                                    --img_dir "test_data/single_images"

The generated results will be saved to the --img_dir.

Fitting a Single Image

# Fitting a single image using HeadNeRF
python FittingSingleImage.py --model_path "TrainedModels/model_Reso32HR.pth" \
                             --img "test_data/single_images/img_000037.png" \
                             --mask "test_data/single_images/img_000037_mask.png" \
                             --para_3dmm "test_data/single_images/img_000037_nl3dmm.pkl" \
                             --save_root "test_data/fitting_res" \
                             --target_embedding "LatentCodeSamples/*/S025_E14_I01_P02.pth"

Args:

Results:

Note:

Citation

If you find our work useful in your research, please consider citing our paper:

@inproceedings{hong2021headnerf,
     author     = {Yang Hong and Bo Peng and Haiyao Xiao and Ligang Liu and Juyong Zhang},
     title      = {HeadNeRF: A Real-time NeRF-based Parametric Head Model},
     booktitle  = {{IEEE/CVF} Conference on Computer Vision and Pattern Recognition (CVPR)},
     year       = {2022}
  }

If you have questions, please contact hymath@mail.ustc.edu.cn.

Acknowledgments

License

Academic or non-profit organization noncommercial research use only.

<!-- - This research was supported by the National Key R&D Program of China (2020YFC1523102), the National Natural Science Foundation of China (No.62122071, 62025207), the Youth Innovation Promotion Association CAS (No. 2018495) and the Fundamental Research Funds for the Central Universities (No.WK3470000021). -->