Putting NeRF on a Diet: Semantically Consistent Few-Shot View Synthesis Implementation
This project attempted to implement the paper Putting NeRF on a Diet (DietNeRF) in JAX/Flax. DietNeRF is designed for rendering quality novel views in few-shot learning scheme, a task that vanilla NeRF (Neural Radiance Field) struggles. To achieve this, the author coins Semantic Consistency Loss to supervise DietNeRF by prior knowledge from CLIP Vision Transformer. Such supervision enables DietNeRF to learn 3D scene reconstruction with CLIP's prior knowledge on 2D views.
Besides this repo, you can check our write-up and demo here:
-
✍️ Write-up in Notion: more details of DietNeRF and our experiments -
✨ Demo in Hugging Face Space: showcase our trained DietNeRFs by Streamlit
🤩
Demo
- You can check out our demo in Hugging Face Space
- Or you can set up our Streamlit demo locally (model checkpoints will be fetched automatically upon startup)
pip install -r requirements_demo.txt
streamlit run app.py
✨
Implementation
Our code is written in JAX/ Flax and mainly based upon jaxnerf from Google Research. The base code is highly optimized in GPU & TPU. For semantic consistency loss, we utilize pretrained CLIP Vision Transformer from transformers library.
To learn more about DietNeRF, our experiments and implementation, you are highly recommended to check out our very detailed Notion write-up!
🤗
Hugging Face Model Hub Repo
You can also find our project and our model checkpoints on our Hugging Face Model Hub Repository. The models checkpoints are located in models folder.
Our JAX/Flax implementation currently supports:
| Platform | Single-Host GPU | Multi-Device TPU | ||
|---|---|---|---|---|
| Type | Single-Device | Multi-Device | Single-Host | Multi-Host |
| Training |  |
|
|
|
| Evaluation | |
|
|
|
💻
Installation
# Clone the repo
git clone https://github.com/codestella/putting-nerf-on-a-diet
# Create a conda environment, note you can use python 3.6-3.8 as
# one of the dependencies (TensorFlow) hasn't supported python 3.9 yet.
conda create --name jaxnerf python=3.6.12; conda activate jaxnerf
# Prepare pip
conda install pip; pip install --upgrade pip
# Install requirements
pip install -r requirements.txt
# [Optional] Install GPU and TPU support for Jax
# Remember to change cuda101 to your CUDA version, e.g. cuda110 for CUDA 11.0.
!pip install --upgrade jax "jax[cuda110]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
# install flax and flax-transformer
pip install flax transformers[flax]
⚽
Dataset
Download the datasets from the NeRF official Google Drive. Please download the nerf_synthetic.zip and unzip them in the place you like. Let's assume they are placed under /tmp/jaxnerf/data/.
🤟
How to Train
- Train in our prepared Colab notebook: Colab Pro is recommended, otherwise you may encounter out-of-memory
- Train locally: set
use_semantic_loss=truein youryamlconfiguration file to enable DietNeRF.
python -m train \
--data_dir=/PATH/TO/YOUR/SCENE/DATA \ # (e.g. nerf_synthetic/lego)
--train_dir=/PATH/TO/THE/PLACE/YOU/WANT/TO/SAVE/CHECKPOINTS \
--config=configs/CONFIG_YOU_LIKE
💎
Experimental Results
❗
Rendered Rendering images by 8-shot learned DietNeRF
DietNeRF has a strong capacity to generalise on novel and challenging views with EXTREMELY SMALL TRAINING SAMPLES!
HOTDOG / DRUM / SHIP / CHAIR / LEGO / MIC
 |
 |
 |
 |
 |
 |
❗
Rendered GIF by occluded 14-shot learned NeRF and Diet-NeRF
We made artificial occlusion on the right side of image (Only picked left side training poses). The reconstruction quality can be compared with this experiment. DietNeRF shows better quality than Original NeRF when It is occluded.
Training poses
LEGO
| Diet NeRF | NeRF |
 |
 |
SHIP
| Diet NeRF | NeRF |
 |
 |
👨👧👦
Our Team
| Teams | Members |
|---|---|
| Project Managing | Stella Yang To Watch Our Project Progress, Please Check Our Project Notion |
| NeRF Team | Stella Yang, Alex Lau, Seunghyun Lee, Hyunkyu Kim, Haswanth Aekula, JaeYoung Chung |
| CLIP Team | Seunghyun Lee, Sasikanth Kotti, Khalid Sifullah , Sunghyun Kim |
| Cloud TPU Team | Alex Lau, Aswin Pyakurel, JaeYoung Chung, Sunghyun Kim |
*Special mention to our "night owl" contributors
💞
Social Impact
- Game Industry
- Augmented Reality Industry
- Virtual Reality Industry
- Graphics Industry
- Online shopping
- Metaverse
- Digital Twin
- Mapping / SLAM
🌱
References
This project is based on “JAX-NeRF”.
@software{jaxnerf2020github,
author = {Boyang Deng and Jonathan T. Barron and Pratul P. Srinivasan},
title = {{JaxNeRF}: an efficient {JAX} implementation of {NeRF}},
url = {https://github.com/google-research/google-research/tree/master/jaxnerf},
version = {0.0},
year = {2020},
}
This project is based on “Putting NeRF on a Diet”.
@misc{jain2021putting,
title={Putting NeRF on a Diet: Semantically Consistent Few-Shot View Synthesis},
author={Ajay Jain and Matthew Tancik and Pieter Abbeel},
year={2021},
eprint={2104.00677},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
🔑
License
❤️
Special Thanks
Our Project is motivated by HuggingFace X GoogleAI (JAX) Community Week Event 2021.
We would like to take this chance to thank Hugging Face for organizing such an amazing open-source initiative, Suraj and Patrick for all the technical help. We learn a lot throughout this wonderful experience!
Finally, we would like to thank Common Computer AI for sponsoring our team access to V100 multi-GPUs server. Thank you so much for your support!
![[on Test] Add data_loader part in trainer.py](https://avatars.githubusercontent.com/u/31179406?v=4)
78 Dec 10, 2022
94 Nov 22, 2022
40 Sep 26, 2022
406 Dec 23, 2022
68 Dec 30, 2022
551 Dec 29, 2022
319 Dec 30, 2022
220 Dec 31, 2022
16 Oct 5, 2022
1 Feb 7, 2022
11 Feb 8, 2022
138 Dec 28, 2022
169 Dec 30, 2022
93 Dec 17, 2022
62 Nov 21, 2022
1 May 24, 2022
126 Dec 30, 2022
186 Dec 30, 2022
585 Jan 4, 2023