yinyunie/Total3DUnderstanding

Total3DUnderstanding [Project Page][Oral Paper][Talk]

Total3DUnderstanding: Joint Layout, Object Pose and Mesh Reconstruction for Indoor Scenes from a Single Image
Yinyu Nie, Xiaoguang Han, Shihui Guo, Yujian Zheng, Jian Chang, Jian Jun Zhang
In CVPR, 2020.

Install

This implementation uses Python 3.6, Pytorch1.1.0, cudatoolkit 9.0. We recommend to use conda to deploy the environment.

Install with conda:

conda env create -f environment.yml
conda activate Total3D

Install with pip:

pip install -r requirements.txt

Demo

The pretrained model can be download here. We also provide the pretrained Mesh Generation Net here. Put the pretrained models under

out/pretrained_models

A demo is illustrated below to see how the method works. vtk is used here to visualize the 3D scenes. The outputs will be saved under 'demo/outputs'. You can also play with your toy with this script.

cd Total3DUnderstanding
python main.py configs/total3d.yaml --mode demo --demo_path demo/inputs/1

Data preparation

In our paper, we use SUN-RGBD to train our Layout Estimation Net (LEN) and Object Detection Net (ODN), and use Pix3D to train our Mesh Generation Net (MGN).

Preprocess SUN-RGBD data

You can either directly download the processed training/testing data [link] to (recommended)

data/sunrgbd/sunrgbd_train_test_data

Download the raw SUN-RGBD data to

data/sunrgbd/Dataset/SUNRGBD

Download the 37 class labels of objects in SUN RGB-D images [link] to

data/sunrgbd/Dataset/SUNRGBD/train_test_labels

Follow this work to download the preprocessed clean data of SUN RGB-D [link] to

'data/sunrgbd/Dataset/data_clean'

Follow this work to download the preprocessed ground-truth of SUN RGB-D [link], and put the '3dlayout' and 'updated_rtilt' folders respectively to

data/sunrgbd/Dataset/3dlayout
data/sunrgbd/Dataset/updated_rtilt

Run below to generate training and testing data in 'data/sunrgbd/sunrgbd_train_test_data'.

python utils/generate_data.py

If everything goes smooth, a ground-truth scene will be visualized like

Preprocess Pix3D data

You can either directly download the preprocessed ground-truth data [link] to (recommended)

data/pix3d/train_test_data

Each sample contains the object class, 3D points (sampled on meshes), sample id and object image (w.o. mask). Samples in the training set are flipped for augmentation.

Download the Pix3D dataset to

data/pix3d/metadata

Run below to generate the train/test data into 'data/pix3d/train_test_data'

python utils/preprocess_pix3d.py

Training and Testing

We use the configuration file (see 'configs/****.yaml') to fully control the training/testing process. There are three subtasks in Total3D (layout estimation, object detection and mesh reconstruction). We first pretrain each task individually followed with joint training.

Pretraining

Switch the keyword in 'configs/total3d.yaml' between ('layout_estimation', 'object_detection') as below to pretrain the two tasks individually.

train:
  phase: 'layout_estimation' # or 'object_detection'

python main.py configs/total3d.yaml --mode train

The two pretrained models can be correspondingly found at

out/total3d/a_folder_named_with_script_time/model_best.pth

Train the Mesh Generation Net by:

python main.py configs/mgnet.yaml --mode train

The pretrained model can be found at

out/mesh_gen/a_folder_named_with_script_time/model_best.pth

Joint training

List the addresses of the three pretrained models in 'configs/total3d.yaml', and modify the phase name to 'joint' as

weight: ['folder_to_layout_estimation/model_best.pth', 'folder_to_object_detection/model_best.pth', 'folder_to_mesh_recon/model_best.pth']

train:
  phase: 'joint'

Then run below for joint training.

python main.py configs/total3d.yaml --mode train

The trained model can be found at

out/total3d/a_folder_named_with_script_time/model_best.pth

Testing

Please make sure the weight path is renewed as

weight: ['folder_to_fully_trained_model/model_best.pth']

and run

python main.py configs/total3d.yaml --mode test

This script generates all 3D scenes on the test set of SUN-RGBD under

out/total3d/a_folder_named_with_script_time/visualization

You can also visualize a 3D scene given the sample id as

python utils/visualize.py --result_path out/total3d/a_folder_named_with_script_time/visualization --sequence_id 274

Differences to the paper

We retrained the model with the learning rate decreases to half if there is no gain within five steps, which is much more efficient.
We do not provide the Faster RCNN code. Users can train their 2D detector with [link].

Citation

If you find our work is helpful, please cite

@InProceedings{Nie_2020_CVPR,
author = {Nie, Yinyu and Han, Xiaoguang and Guo, Shihui and Zheng, Yujian and Chang, Jian and Zhang, Jian Jun},
title = {Total3DUnderstanding: Joint Layout, Object Pose and Mesh Reconstruction for Indoor Scenes From a Single Image},
booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
month = {June},
year = {2020}
}

Our method partially follows the data processing steps in this work. If it is also helpful to you, please cite

@inproceedings{huang2018cooperative,
  title={Cooperative Holistic Scene Understanding: Unifying 3D Object, Layout, and Camera Pose Estimation},
  author={Huang, Siyuan and Qi, Siyuan and Xiao, Yinxue and Zhu, Yixin and Wu, Ying Nian and Zhu, Song-Chun},
  booktitle={Advances in Neural Information Processing Systems},
  pages={206--217},
  year={2018}
}

yinyunie / Total3DUnderstanding