/keras_nerf

NeRF (Neural Radiance Field) TensorFlow v2 Keras Re-Implementation

Primary LanguagePython

NeRF TensorFlow v2 Keras Re-Implementation

[Work in Progress]

Author: Naufal Suryanto

Quickstart

1. Prepare dataset

  • Download the nerf_synthetic dataset from NeRF Official [Here]
  • Extract the dataset to data folder. The recommended structure is data/nerf_synthetic/[object_name]

2. Install the dependencies

  • Make sure that you have install all dependencies
  • To install the dependencies, you can run python3 -m pip install -r requirements.txt

3. Run the tests [optional]

  • Run the tests to make sure that all functions can be run as expected python3 -m pytest .

4. Run the training code

  • Single GPU: If you only have a single gpu, run the training code train_single.py
  • Multiple GPUs: If you only have multiple gpu, run the training code train.py. The batch size will be multiplied by number of GPU instances
  • Examples:
    • python3 train.py --data_dir data/nerf_synthetic/lego --img_wh 128 --batch_size 2 --name lego_128 --ray_chunks 2048 --white_bg
    • python3 train_single.py --data_dir data/nerf_synthetic/chair --img_wh 128 --eagerly --name chair_128 --ray_chunks 4096
  • By default, the logs and trained model can be seen on logs and models directories, respectively

5. Run the inference code

  • After complete training the model, we can run the inference code to generate 360 degrees rotation of the model
  • Run the inference code, for example: python3 inference.py --model_dirs model/lego_128 --img_wh 256 --name lego_128_256
  • By default, the output gif can be seen on output directory

Results

Training Logs Inference
lego_train_250ep lego_128_256
chair_train_245ep chair_128_256

Implementation and Features

NeRF Model

  • Positional Encoding
  • Coarse and Fine Model with Hierarchical Sampling

Training Supports

  • Single or multiple GPU training
  • Split the MLP prediction into chunks for fitting GPU memory (tune --ray_chunks argument to fit your memory; should be a factors of image_height * image_width * batch_size)
  • [Default] Use graph execution with tf.function for better performance (but may take longer initialization and larger memory usage)
  • [Optional] Use eager execution for faster initialization and lower memory usage (include --eagerly option when run the code, but currenly only works on single GPU training)
  • Black or White Background (include --white_bg option when run the code for better result)
  • Log the training history in CSV, plot image, and sample image
  • Continue model training from last logging step

References

Code Implementation Inspired by

Original Project Page