[Paper
]
$\rm O^2$ -Recon: Completing 3D Reconstruction of Occluded Objects in the Scene with a Pre-trained 2D Diffusion Model
Yubin Hu, Sheng Ye, Wang Zhao, Matthieu Lin, Yuze He, Yu-Hui Wen, Ying He, Yong-Jin Liu
AAAI 2024
We will release the full code as soon as possible.
- Release the code of training on an example object data.
- Release the dataset pre-processing code.
- Update readme.
Install the required packages:
conda create -n O2-recon python=3.7
conda activate O2-recon
pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
git clone --recurse-submodules https://github.com/THU-LYJ-Lab/O2-Recon.git
cd O2-Recon
pip install -r requirements.txt
pip install git+https://github.com/openai/CLIP.git
Download the normal prediction model scannet_neuris_retrain.pt
from folder jiepeng@connect.hku.hk > GitHub >NeuRlS > pretrained normal network > snu
in this OneDrive and store it to ./preprocess/surface_normal_uncertainty/checkpoints/
You can download the in-painted dataset from here: .
After downloaded, unzip the file into the directory ./dataset/indoor-paper
.
We provide an example to prepare the dataset from scannet format.
Download several scenes of the whole scannet dataset. You can download the scenes we utilized from here. Unzip the files into ./scannet/scenexxxx_xx_scannet
directories.
Depending on which scenes you'd like to process, you need to modify L150 in preprocess/object_mask_with_clip.py
. And then run
python preprocess/object_mask_with_clip.py
This script extracts objects from the scannet scenes according to instance masks and semantic categories. As a result, the object data is extracted into ./scannet/object_original_with_clip/scenexxxx_xx_scannet_obj_x
.
You can download our created masks (TsinghuaCloud, GoogleDrive). Or you can generate the masks and name them following the same manner. We utilize the xx_class_inpaint mask.png
file under each directory.
After downloaded, place the directories to correct locations. For example, place the only-seg-0008-obj10
directory under /path/to/O2-Recon/scannet/object_original_with_clip/scene0008_00_scannet_obj_10/
.
Based on these object directories, next we generate the in-paintings and predict the monocular cues.
Depending on which scenes you'd like to process, you need to modify L41 in exp_preprocess.py
.
And then run
python exp_preprocess.py --data_type scannet-with-inpaint --scannet_root=/path/to/O2-Recon/scannet/ --neus_root=/path/to/O2-Recon/dataset/indoor-paper/ --dir_snu_code /path/to/O2-Recon/preprocess/surface_normal_uncertainty/
Here you need to use absolute paths.
Run training and mesh extraction with our example script. For example, run the following script to reconstruct all objects in scene scene0005_00
:
bash scripts/train_0005.sh
This script reconstruct the objects one-by-one. After all the steps finished, the reconstructed results and middle results can be found under /path/to/O2-Recon/exps-paper/indoor-paper/neus
.