Putting NeRF on a Diet: Semantically Consistent Few-Shot View Synthesis Implementation
the Pytorch, JAX/Flax based code imlementation of this paper : https://arxiv.org/abs/2104.00677 The model gives the 3D neural scene representation (NeRF: Neural Radiances Field) estimated from a few images. Which is based on extracting the semantic information using a pre-trained visual encoder such as CLIP, a Vision Transformer
Our Project is started in the HuggingFace X GoogleAI (JAX) Community Week Event. https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/7104
Hugging Face Hub Repo URL:
We will also upload our project on the Hugging Face Hub Repository. https://huggingface.co/flax-community/putting-nerf-on-a-diet/
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
svn export https://github.com/google-research/google-research/trunk/jaxnerf
# 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 jaxnerf/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 jaxlib==0.1.57+cuda101 -f https://storage.googleapis.com/jax-releases/jax_releases.html
# install flax and flax-transformer
pip install flax transformer[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 use
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
you can toggle the semantic loss by “use_semantic_loss” in cofiguration files.
Rendered examples by 8-shot learned Diet-NeRF
- Lego
- Chair
Rendered examples by occluded 14-shot learned NeRF and Diet-NeRF
This result is on the quite initial state and expected to be improved.
Training poses
Rendered novel poses
Demo
https://huggingface.co/spaces/flax-community/DietNerf-Demo
Our Teams
Teams | Members |
---|---|
NeRF Team | Leader : JaeYoung Chung, Members : Stella Yang, Alex Lau, Haswanth Aekula, Hyunkyu Kim |
CLIP Team | Leader : Sunghyun Kim, Members : Seunghyun Lee, Sasikanth Kotti, Khali Sifullah |
Cloud TPU Team | Leader : Alex Lau, Members : JaeYoung Chung, Sunghyun Kim, Aswin Pyakurel |
Project Managing | Stella Yang To Watch Our Project Progress, Please Check Our Project Notion |
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 “JAX-NeRF”.
@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}
}