Haian-Jin / TensoIR

[CVPR 2023] TensoIR: Tensorial Inverse Rendering
https://haian-jin.github.io/TensoIR/
MIT License
237 stars 12 forks source link

TensoIR: Tensorial Inverse Rendering (CVPR 2023)

Project Page | Paper

This repository contains a pytorch implementation for the paper: TensoIR: Tensorial Inverse Rendering.

The code can run well, but it is not well organized. I may re-organize the code when I am available.

https://user-images.githubusercontent.com/79512936/235218355-0d4177c1-7614-4772-a8ec-44d76a95743f.mp4

Tested on Ubuntu 20.04 + Pytorch 1.10.1

Install environment:

conda create -n TensoIR python=3.8
conda activate TensoIR
pip install torch==1.10 torchvision
pip install tqdm scikit-image opencv-python configargparse lpips imageio-ffmpeg kornia lpips tensorboard loguru plyfile

Dataset

Downloading

Please download the dataset and environment maps from the following links and put them in the ./data folder:

Generating your own synthetic dataset

We provide the code for generating your own synthetic dataset with your own Blender files and Blender software. Please download this file and follow the readme.md file inside it to render your own dataset. The Blender rendering scripts heavily rely on the code provided by NeRFactor. Thanks for its great work!

Training

Note:

  1. After finishing all training iterations, the training script will automatically render all test images under the learned lighting condition and save them in the log folder. It will also compute all metrics related to geometry, materials, and novel view synthesis(except for relighting). The results will be saved in the log folder as well.
  2. Different scenes have different config files. The main difference between those config files is the different weight values for normals_diff_weight, which controls how close the predicted normals should be to the derived normals. A larger weight will help prevent the normals prediction from overfitting the supervised colors, but at the same time, it will damage the normals prediction network's ability to predict high-frequency details. We recommend three values to try: 0.0005, 0.002, and 0.005 when you train TensoIR on your own dataset.

For pre-trained checkpoints and results please see:

Checkpoints Results

Training under single lighting condition

export PYTHONPATH=. && python train_tensoIR.py --config ./configs/single_light/armadillo.txt

Training under rotated multi-lighting conditions

export PYTHONPATH=. && python train_tensoIR_rotated_multi_lights.py  --config ./configs/multi_light_rotated/hotdog.txt

Training under general multi-lighting conditions

export PYTHONPATH=. && python train_tensoIR_general_multi_lights.py  --config ./configs/multi_light_general/ficus.txt

(Optional) Training for the original NeRF-Synthetic dataset

We don't do quantitative and qualitative comparisons for the original NeRF-Synthetic dataset in our paper (the reasons have been discussed above), but you can still train TensoIR on the original NeRF-Synthetic dataset for some analysis.

export PYTHONPATH=. && python train_tensoIR_simple.py --config ./configs/single_light/blender.txt

Testing and Validation

Rendering with a pre-trained model under learned lighting condition

export PYTHONPATH=. && python "$training_file" --config "$config_path" --ckpt "$ckpt_path" --render_only 1 --render_test 1

"$training_file" is the training script you used for training, e.g. train_tensoIR.py or train_tensoIR_rotated_multi_lights.py or train_tensoIR_general_multi_lights.py.

"$config_path" is the path to the config file you used for training, e.g. ./configs/single_light/armadillo.txt or ./configs/multi_light_rotated/hotdog.txt or ./configs/multi_light_general/ficus.txt.

"$ckpt_path" is the path to the checkpoint you want to test.

The result will be stored in --basedir defined in the config file.

Relighting with a pre-trained model under unseen lighting conditions

export PYTHONPATH=. && python scripts/relight_importance.py --ckpt "$ckpt_path" --config configs/relighting_test/"$scene".txt --batch_size 800

We do light-intensity importance sampling for relighting. The sampling results are stored in --geo_buffer_path defined in the config file.

"$ckpt_path" is the path to the checkpoint you want to test.

"$scene" is the name of the scene you want to relight, e.g. armadillo or ficus or hotdog or lego.

Reduce the --batch_size if you have limited GPU memory.

The line 370 of scripts/relight_importance.py specifies the names of environment maps for relighting. You can change it if you want to test other unseen lighting conditions.

Extracting mesh

The mesh will be stored in the same folder as the checkpoint.

export PYTHONPATH=. && python scripts/export_mesh.py --ckpt "$ckpt_path" 

Citations

If you find our code or paper helps, please consider citing:

@inproceedings{Jin2023TensoIR,
  title={TensoIR: Tensorial Inverse Rendering},
  author={Jin, Haian and Liu, Isabella and Xu, Peijia and Zhang, Xiaoshuai and Han, Songfang and Bi, Sai and Zhou, Xiaowei and Xu, Zexiang and Su, Hao},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2023}
}

Acknowledgement

The code was built on TensoRF. Thanks for this great project!