SGAM: Building a Virtual 3D World through Simultaneous Generation and Mapping

Yuan Shen¹, Wei-Chiu Ma², Shenlong Wang¹
Unversity of Illinois at Urbana-Champaign¹, Massachusetts Institute of Technology²

Accepted at NeurIPS 2022.

Paper link │ Project Page │ Colab Quickstart

(The GIF animation above is generated via SGAM with only the first RGB-D frame known.)

TL;DR

We present a new 3D scene generation framework that simultaneously generates sensor data at novel viewpoints and builds a 3D map. Our framework is illustrated in the diagram below.

Quickstart

Try our Colab codebook to play our trained models on CLEVR-Infinite and GoogleEarth-Infinite!

Installment

• Manual Installment (Only tested on Ubuntu 18.04):

  1. Create Conda environment and install part of python packages

     conda create -n sgam python=3.9.13
     conda activate sgam
     pip install -r requirement.txt
     

  2. Install pytorch and pytorch_lightning

     pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113
     pip install pytorch_lightning==1.5.10
     

Video

Data

CLEVR-Infinite Dataset

• Note the depth is nonlinear when we render from blender. Checkout how we convert the depth from nonlinear to linear in line 103 of data/clevr-infinite.py
• To get a quick glance at our dataset, here is one tiny scene example.
• Two validation scene data can be downloaded from this link.
• Our training, validation and testing dataset is available at this link for downloading.
• To generate more training dataset at a large scale, we provide the blender script in clevr_generation directory. We randomly distribute primitive 3D objects by simulating flying objects falling and collision. Detailed steps are as follows:
  1. find a device that has GPU, and then install blender 2.92

  sudo snap install blender --channel=2.92/stable --classic
  

  2. (optional) If you want to visualize one CLEVR-Infinite scene, run the following command.

  /snap/bin/blender random_scene.blend
  

  3. Specify output directory in line 253

  4. Run the following command to render. You can change the iteration number to set the number of random scene.

    bash blender_generation.sh
  

  5. Run the postprocessing script to get rgb, depth map and transform.json

    python convert_exr.py
  

GoogleEarth-Infinite Dataset

• Train, val and test set are available at this google drive link.

Trained Models

We provide our trained model on GoogleEarth-Infinite and CLEVR-Infinite. Please download and organize the pre-trained checkpoints as follows:

SGAM   
└───trained_models
    └───google_earth
    │   │   config.yaml
    │   │   XXX.ckpt
    │    
    └───clevr-infinite   
        │   config.yaml
        │   XXX.ckpt

Training

1. VQGAN codebook training.

python train_generative_sensing_model.py --base configs/codebooks/XXX.yaml --gpus 0, -t True

2. Conditional Generation

python train_generative_sensing_model.py --base configs/conditional_generation/XXX.yaml --gpus 0, -t True

Inference

CLEVR-Infinite

python main_scene_generation.py --dataset="clevr-infinite" --use_rgbd_integration True

GoogleEarth-Infinite

python main_scene_generation.py --dataset="google_earth" --use_rgbd_integration True

Acknowledgement

We thank Vlas Zyrianov for his feedback on our paper drafts. Besides, our codebase is modified on top of VQGAN codebase. Many thanks to Patrick Esser and Robin Rombach, who makes their code available.

Citation

If you find our work is useful, please cite our work with the bibtex down below, thanks!

@inproceedings{
    shen2022sgam,
    title={{SGAM}: Building a Virtual 3D World through Simultaneous Generation and Mapping},
    author={Yuan Shen and Wei-Chiu Ma and Shenlong Wang},
    booktitle={Thirty-Sixth Conference on Neural Information Processing Systems},
    year={2022},
    url={https://openreview.net/forum?id=17KCLTbRymw}
}