"SinNeRF: Training Neural Radiance Fields on Complex Scenes from a Single Image", Dejia Xu, Yifan Jiang, Peihao Wang, Zhiwen Fan, Humphrey Shi, Zhangyang Wang

SinNeRF: Training Neural Radiance Fields on Complex Scenes from a Single Image

[Paper] [Website]

Pipeline

Code

Environment

pip install -r requirements.txt

Dataset Preparation

Please download the datasets from these links:

Please download the depth from here: https://drive.google.com/drive/folders/13Lc79Ox0k9Ih2o0Y9e_g_ky41Nx40eJw?usp=sharing

Training

If you meet OOM issue, try:

  1. enable precision=16
  2. reduce the patch size --patch_size (or --patch_size_x, --patch_size_y) and enlarge the stride size --sH, --sW
NeRF synthetic
  • Step 1

    python train.py  --dataset_name blender_ray_patch_1image_rot3d  --root_dir  ../../dataset/nerf_synthetic/lego   --N_importance 64 --img_wh 400 400 --num_epochs 3000 --batch_size 1  --optimizer adam --lr 2e-4  --lr_scheduler steplr --decay_step 1000 2000 --decay_gamma 0.5  --exp_name lego_s6 --with_ref --patch_size 64 --sW 6 --sH 6 --proj_weight 1 --depth_smooth_weight 0  --dis_weight 0 --num_gpus 4 --load_depth --depth_type nerf --model sinnerf --depth_weight 8 --vit_weight 10 --scan 4
    
  • Step 2

    python train.py  --dataset_name blender_ray_patch_1image_rot3d  --root_dir  ../../dataset/nerf_synthetic/lego   --N_importance 64 --img_wh 400 400 --num_epochs 3000 --batch_size 1  --optimizer adam --lr 1e-4  --lr_scheduler steplr --decay_step 1000 2000 --decay_gamma 0.5  --exp_name lego_s6_4ft --with_ref --patch_size 64 --sW 4 --sH 4 --proj_weight 1 --depth_smooth_weight 0.1  --dis_weight 0.1 --num_gpus 4 --load_depth --depth_type nerf --model sinnerf --depth_weight 8 --vit_weight 0 --pt_model xxx.ckpt --nerf_only  --scan 4
    
LLFF
  • Step 1

    python train.py  --dataset_name llff_ray_patch_1image_proj  --root_dir  ../../dataset/nerf_llff_data/room   --N_importance 64 --img_wh 504 378 --num_epochs 3000 --batch_size 1  --optimizer adam --lr 2e-4  --lr_scheduler steplr --decay_step 1000 2000 --decay_gamma 0.5  --exp_name llff_room_s4 --with_ref --patch_size_x 63 --patch_size_y 84 --sW 4 --sH 4 --proj_weight 1 --depth_smooth_weight 0  --dis_weight 0 --num_gpus 4 --load_depth --depth_type nerf --model sinnerf --depth_weight 8 --vit_weight 10
    
  • Step 2

    python train.py  --dataset_name llff_ray_patch_1image_proj  --root_dir  ../../dataset/nerf_llff_data/room   --N_importance 64 --img_wh 504 378 --num_epochs 3000 --batch_size 1  --optimizer adam --lr 1e-4  --lr_scheduler steplr --decay_step 1000 2000 --decay_gamma 0.5  --exp_name llff_room_s4_2ft --with_ref --patch_size_x 63 --patch_size_y 84 --sW 2 --sH 2 --proj_weight 1 --depth_smooth_weight 0.1  --dis_weight 0.1 --num_gpus 4 --load_depth --depth_type nerf --model sinnerf --depth_weight 8 --vit_weight 0 --pt_model xxx.ckpt --nerf_only
    
DTU
  • Step 1

    python train.py  --dataset_name dtu_proj  --root_dir  ../../dataset/mvs_training/dtu   --N_importance 64 --img_wh 640 512 --num_epochs 3000 --batch_size 1  --optimizer adam --lr 2e-4  --lr_scheduler steplr --decay_step 1000 2000 --decay_gamma 0.5  --exp_name dtu_scan4_s8 --with_ref --patch_size_y 70 --patch_size_x 56 --sW 8 --sH 8 --proj_weight 1 --depth_smooth_weight 0  --dis_weight 0 --num_gpus 4 --load_depth --depth_type nerf --model sinnerf --depth_weight 8 --vit_weight 10 --scan 4
    
  • Step 2

    python train.py  --dataset_name dtu_proj  --root_dir  ../../dataset/mvs_training/dtu   --N_importance 64 --img_wh 640 512 --num_epochs 3000 --batch_size 1  --optimizer adam --lr 1e-4  --lr_scheduler steplr --decay_step 1000 2000 --decay_gamma 0.5  --exp_name dtu_scan4_s8_4ft --with_ref --patch_size_y 70 --patch_size_x 56 --sW 4 --sH 4 --proj_weight 1 --depth_smooth_weight 0.1  --dis_weight 0.1 --num_gpus 4 --load_depth --depth_type nerf --model sinnerf --depth_weight 8 --vit_weight 0 --pt_model xxx.ckpt --nerf_only  --scan 4
    

More finetuning with smaller strides benefits reconstruction quality.

Testing

python eval.py  --dataset_name llff  --root_dir /dataset/nerf_llff_data/room --N_importance 64 --img_wh 504 378 --model nerf --ckpt_path ckpts/room.ckpt --timestamp test

Acknowledgement

Codebase based on https://github.com/kwea123/nerf_pl . Thanks for sharing!

Citation

If you find this repo is helpful, please cite:


@InProceedings{Xu_2022_SinNeRF,
author = {Xu, Dejia and Jiang, Yifan and Wang, Peihao and Fan, Zhiwen and Shi, Humphrey and Wang, Zhangyang},
title = {SinNeRF: Training Neural Radiance Fields on Complex Scenes from a Single Image},
journal={arXiv preprint arXiv:2204.00928},
year={2022}
}

Owner
VITA
Visual Informatics Group @ University of Texas at Austin
VITA
Comments
  • problems in load vit

    problems in load vit

    Hi, thank for your exciting work,but When I tried to train in the room scene, I had the following problems when loading the VIT model, could you give me some suggestions?

    ~/NeRFs/SinNeRF$ python train.py --dataset_name llff_ray_patch_1image_proj --root_dir data/nerf_llff_data/room --N_importance 64 --img_wh 504 378 --num_epochs 3000 --batch_size 1 --optimizer adam --lr 2e-4 --lr_scheduler steplr --decay_step 1000 2000 --decay_gamma 0.5 --exp_name llff_room_s4 --with_ref --patch_size_x 63 --patch_size_y 84 --sW 4 --sH 4 --proj_weight 1 --depth_smooth_weight 0 --dis_weight 0 --num_gpus 4 --load_depth --depth_type nerf --model sinnerf --depth_weight 8 --vit_weight 10 Namespace(N_importance=64, N_samples=64, angle=30, batch_size=1, chunk=32768, ckpt_path=None, dataset_name='llff_ray_patch_1image_proj', decay_gamma=0.5, decay_step=[1000, 2000], depth_anneal=False, depth_smooth_weight=0.0, depth_type='nerf', depth_weight=8.0, dis_weight=0.0, dloss='hinge', exp_name='llff_room_s4', img_wh=[504, 378], load_depth=True, loss_type='mse', lr=0.0002, lr_scheduler='steplr', model='sinnerf', momentum=0.9, nH=32, nW=32, nerf_only=False, noise_std=1.0, num_epochs=3000, num_gpus=4, optimizer='adam', patch_loss='mse', patch_size=-1, patch_size_x=63, patch_size_y=84, perturb=1.0, poly_exp=0.9, prefixes_to_ignore=['loss'], proj_weight=1.0, pt_model=None, repeat=1, root_dir='data/nerf_llff_data/room', sH=4, sW=4, scan=4, spheric_poses=False, use_disp=False, vit_weight=10.0, warmup_epochs=0, warmup_multiplier=1.0, weight_decay=0, with_ref=True) Using cache found in /home/zhangzhongwei18/.cache/torch/hub/facebookresearch_dino_main Traceback (most recent call last): File "train.py", line 19, in system = SinNeRF(hparams) File "/home/zhangzhongwei18/NeRFs/SinNeRF/models/sinnerf.py", line 148, in init self.ext = VitExtractor( File "/home/zhangzhongwei18/NeRFs/SinNeRF/models/extractor.py", line 22, in init self.model = torch.hub.load( File "/home/zhangzhongwei18/.custom/cuda-10.2-cudnn8-devel-ubuntu18.04-pytorch1.8.0_full_tensorboard/envs/sinnerf/lib/python3.8/site-packages/torch/hub.py", line 404, in load model = _load_local(repo_or_dir, model, *args, **kwargs) File "/home/zhangzhongwei18/.custom/cuda-10.2-cudnn8-devel-ubuntu18.04-pytorch1.8.0_full_tensorboard/envs/sinnerf/lib/python3.8/site-packages/torch/hub.py", line 430, in _load_local hub_module = _import_module(MODULE_HUBCONF, hubconf_path) File "/home/zhangzhongwei18/.custom/cuda-10.2-cudnn8-devel-ubuntu18.04-pytorch1.8.0_full_tensorboard/envs/sinnerf/lib/python3.8/site-packages/torch/hub.py", line 76, in import_module spec.loader.exec_module(module) File "", line 783, in exec_module File "", line 219, in call_with_frames_removed File "/home/zhangzhongwei18/.cache/torch/hub/facebookresearch_dino_main/hubconf.py", line 17, in import vision_transformer as vits File "/home/zhangzhongwei18/.cache/torch/hub/facebookresearch_dino_main/vision_transformer.py", line 24, in from utils import trunc_normal ImportError: cannot import name 'trunc_normal' from 'utils' (/home/zhangzhongwei18/NeRFs/SinNeRF/utils/init.py)

  • GPU environments

    GPU environments

    Hi I encountered the problem that CUDA out of memory. My gpu spec is GeForce RTX 2080 Ti. And I'm using DDP.

    Could you tell me your gpu spec? Also, I would appreciate it if you could tell me how to make a depth image .npy files.

  • Training time

    Training time

    Hi, Thanks for the great work and for uplaoding the training code.

    Can you share the time it took to train the model room.ckpt and the device it was trained on so that I can have a reference for reproducing/ training for other scenes?

    Thanks

  • About paper results

    About paper results

    Thanks for your great work. I want to know how you train DS-NeRF from a single view. DS-NeRF needs structure-from-motion to aquire 3D sparse points. Thanks in advance! image

  • Confusing results after step-2 training.

    Confusing results after step-2 training.

    Hi! Thanks for your great work and implementation!

    I'm currently trying your code on nerf_synthetic (lego) and dtu (scan4), while the results after the two-stage training are confusing.

    Specifically, the evaluation psnr for lego is 20.6 by the end of step 1, while it drops to 14.9 after step 2 training. The same thing happens to dtu_scan_4, where the evaluation psnr is around 15.0 after step 1 and drops to 11.8 after step 2. The visualization results for lego are as below.

    I'd appreciate it if you could provide some idea on this phenomenon and how to fix this problem. Thank you!

    lego after step 1 002 lego after step 2 010

[Preprint] "Bag of Tricks for Training Deeper Graph Neural Networks A Comprehensive Benchmark Study" by Tianlong Chen*, Kaixiong Zhou*, Keyu Duan, Wenqing Zheng, Peihao Wang, Xia Hu, Zhangyang Wang
[Preprint]

Bag of Tricks for Training Deeper Graph Neural Networks: A Comprehensive Benchmark Study Codes for [Preprint] Bag of Tricks for Training Deeper Graph

Jun 25, 2022
[ICML 2021] “ Self-Damaging Contrastive Learning”, Ziyu Jiang, Tianlong Chen, Bobak Mortazavi, Zhangyang Wang
[ICML 2021] “ Self-Damaging Contrastive Learning”, Ziyu Jiang, Tianlong Chen, Bobak Mortazavi, Zhangyang Wang

Self-Damaging Contrastive Learning Introduction The recent breakthrough achieved by contrastive learning accelerates the pace for deploying unsupervis

Jun 30, 2022
[ICLR 2021] "Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective" by Wuyang Chen, Xinyu Gong, Zhangyang Wang

Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective [PDF] Wuyang Chen, Xinyu Gong, Zhangyang Wang In ICLR 2

Jun 27, 2022
[ICLR 2021 Spotlight Oral] "Undistillable: Making A Nasty Teacher That CANNOT teach students", Haoyu Ma, Tianlong Chen, Ting-Kuei Hu, Chenyu You, Xiaohui Xie, Zhangyang Wang

Undistillable: Making A Nasty Teacher That CANNOT teach students "Undistillable: Making A Nasty Teacher That CANNOT teach students" Haoyu Ma, Tianlong

Jul 3, 2022
[CVPRW 21] "BNN - BN = ? Training Binary Neural Networks without Batch Normalization", Tianlong Chen, Zhenyu Zhang, Xu Ouyang, Zechun Liu, Zhiqiang Shen, Zhangyang Wang
[CVPRW 21]

BNN - BN = ? Training Binary Neural Networks without Batch Normalization Codes for this paper BNN - BN = ? Training Binary Neural Networks without Bat

Apr 28, 2022
[CVPR 2021] "The Lottery Tickets Hypothesis for Supervised and Self-supervised Pre-training in Computer Vision Models" Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Michael Carbin, Zhangyang Wang
[CVPR 2021]

The Lottery Tickets Hypothesis for Supervised and Self-supervised Pre-training in Computer Vision Models Codes for this paper The Lottery Tickets Hypo

Jun 20, 2022
[ICML 2021] "Graph Contrastive Learning Automated" by Yuning You, Tianlong Chen, Yang Shen, Zhangyang Wang
[ICML 2021]

Graph Contrastive Learning Automated PyTorch implementation for Graph Contrastive Learning Automated [talk] [poster] [appendix] Yuning You, Tianlong C

Jul 5, 2022
[Preprint] "Chasing Sparsity in Vision Transformers: An End-to-End Exploration" by Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, Zhangyang Wang
[Preprint]

Chasing Sparsity in Vision Transformers: An End-to-End Exploration Codes for [Preprint] Chasing Sparsity in Vision Transformers: An End-to-End Explora

Jun 21, 2022
[CVPR 2022] "The Principle of Diversity: Training Stronger Vision Transformers Calls for Reducing All Levels of Redundancy" by Tianlong Chen, Zhenyu Zhang, Yu Cheng, Ahmed Awadallah, Zhangyang Wang
[CVPR 2022]

The Principle of Diversity: Training Stronger Vision Transformers Calls for Reducing All Levels of Redundancy Codes for this paper: [CVPR 2022] The Pr

Mar 30, 2022
PyTorch implementation of Super SloMo by Jiang et al.
PyTorch implementation of Super SloMo by Jiang et al.

Super-SloMo PyTorch implementation of "Super SloMo: High Quality Estimation of Multiple Intermediate Frames for Video Interpolation" by Jiang H., Sun

Jun 23, 2022
This is the official PyTorch implementation of the paper "TransFG: A Transformer Architecture for Fine-grained Recognition" (Ju He, Jie-Neng Chen, Shuai Liu, Adam Kortylewski, Cheng Yang, Yutong Bai, Changhu Wang, Alan Yuille).
This is the official PyTorch implementation of the paper

TransFG: A Transformer Architecture for Fine-grained Recognition Official PyTorch code for the paper: TransFG: A Transformer Architecture for Fine-gra

Jul 4, 2022
Code for the ICML 2021 paper "Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Training and Effective Adaptation", Haoxiang Wang, Han Zhao, Bo Li.
Code for the ICML 2021 paper

Bridging Multi-Task Learning and Meta-Learning Code for the ICML 2021 paper "Bridging Multi-Task Learning and Meta-Learning: Towards Efficient Trainin

Jun 10, 2022
[ICCV'2021] "SSH: A Self-Supervised Framework for Image Harmonization", Yifan Jiang, He Zhang, Jianming Zhang, Yilin Wang, Zhe Lin, Kalyan Sunkavalli, Simon Chen, Sohrab Amirghodsi, Sarah Kong, Zhangyang Wang
[ICCV'2021]

SSH: A Self-Supervised Framework for Image Harmonization (ICCV 2021) code for SSH Representative Examples Main Pipeline RealHM DataSet Google Drive Pr

Jun 13, 2022
[Preprint] "Bag of Tricks for Training Deeper Graph Neural Networks A Comprehensive Benchmark Study" by Tianlong Chen*, Kaixiong Zhou*, Keyu Duan, Wenqing Zheng, Peihao Wang, Xia Hu, Zhangyang Wang
[Preprint]

Bag of Tricks for Training Deeper Graph Neural Networks: A Comprehensive Benchmark Study Codes for [Preprint] Bag of Tricks for Training Deeper Graph

Jun 25, 2022
[ICML 2021] “ Self-Damaging Contrastive Learning”, Ziyu Jiang, Tianlong Chen, Bobak Mortazavi, Zhangyang Wang
[ICML 2021] “ Self-Damaging Contrastive Learning”, Ziyu Jiang, Tianlong Chen, Bobak Mortazavi, Zhangyang Wang

Self-Damaging Contrastive Learning Introduction The recent breakthrough achieved by contrastive learning accelerates the pace for deploying unsupervis

Jun 30, 2022
[ICLR 2021] "Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective" by Wuyang Chen, Xinyu Gong, Zhangyang Wang

Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective [PDF] Wuyang Chen, Xinyu Gong, Zhangyang Wang In ICLR 2

Jun 27, 2022
[ICLR 2021 Spotlight Oral] "Undistillable: Making A Nasty Teacher That CANNOT teach students", Haoyu Ma, Tianlong Chen, Ting-Kuei Hu, Chenyu You, Xiaohui Xie, Zhangyang Wang

Undistillable: Making A Nasty Teacher That CANNOT teach students "Undistillable: Making A Nasty Teacher That CANNOT teach students" Haoyu Ma, Tianlong

Jul 3, 2022
[CVPRW 21] "BNN - BN = ? Training Binary Neural Networks without Batch Normalization", Tianlong Chen, Zhenyu Zhang, Xu Ouyang, Zechun Liu, Zhiqiang Shen, Zhangyang Wang
[CVPRW 21]

BNN - BN = ? Training Binary Neural Networks without Batch Normalization Codes for this paper BNN - BN = ? Training Binary Neural Networks without Bat

Apr 28, 2022
[CVPR 2021] "The Lottery Tickets Hypothesis for Supervised and Self-supervised Pre-training in Computer Vision Models" Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Michael Carbin, Zhangyang Wang
[CVPR 2021]

The Lottery Tickets Hypothesis for Supervised and Self-supervised Pre-training in Computer Vision Models Codes for this paper The Lottery Tickets Hypo

Jun 20, 2022
[ICML 2021] "Graph Contrastive Learning Automated" by Yuning You, Tianlong Chen, Yang Shen, Zhangyang Wang
[ICML 2021]

Graph Contrastive Learning Automated PyTorch implementation for Graph Contrastive Learning Automated [talk] [poster] [appendix] Yuning You, Tianlong C

Jul 5, 2022