ODIN: A Single Model for 2D and 3D Segmentation

Authors: Ayush Jain, Pushkal Katara, Nikolaos Gkanatsios, Adam W. Harley, Gabriel Sarch, Kriti Aggarwal, Vishrav Chaudhary, Katerina Fragkiadaki.

Official implementation of "ODIN: A Single Model for 2D and 3D Segmentation", CVPR 2024 (Highlight).

Installation

Make sure you are using a newer version of GCC>=9.2.0

export TORCH_CUDA_ARCH_LIST="6.0 6.1 6.2 7.0 7.2 7.5 8.0 8.6"
conda create -n odin python=3.10
conda activate odin
pip install torch==2.2.0+cu118 torchvision==0.17.0+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
pip install torch-scatter -f https://data.pyg.org/whl/torch-2.2.0+cu118.html
pip install -r requirements.txt
sh init.sh

Data Preparation

Please refer to README in data_preparation folder for individual datasets. For eg. ScanNet data preparation README

Usage

We provide training scripts for various datasets in scripts folder. Please refer to these scripts for training ODIN.

• Modify DETECTRON2_DATASETS to the path where you store the Posed RGB-D data. You might also need to change 3D Mesh point cloud paths (like SCANNET_DATA_DIR) for each script. You may want to find these variables in odin/config.py and permanently modify these paths.
• If you want to train a model on 3D datasets, modify MODEL.WEIGHTS to load pre-trained 2D weights. We use weights from Mask2Former. You can download the Mask2Former-ResNet and Mask2Former-Swin Weights. MODEL.WEIGHTS can also accept link to the checkpoint as well, so you can directly supply these links as the argument value.
• ODIN Pre-trained weights are provided below in the Model-Zoo. Simply point to these weights using the MODEL.WEIGHTS to run inference. You would also need to add --eval-only flag for running evaluation.
• SOLVER.IMS_PER_BATCH controls the batch size. This is effective batch size i.e. if you are running on 2 GPUs and the batch size is set to 6, you are using bs=3 per GPU.
• SOLVER.TEST_IMS_PER_BATCH controls the (effective) test batch size. Since, there are variable number of images in a scene, we use bs=1 per GPU at test time. MAX_FRAME_NUM=-1 means that it loads all images in a scene for inference, which is our usual strategy. In some datasets, the images can simply be too large, thus there we actually set a maximum limit on images.
• INPUT.SAMPLING_FRAME_NUM controls the number of images we sample at test time -- for eg. in ScanNet, we train on 25 image chunks at training time.
• CHECKPOINT_PERIOD is the number of iterations after which a checkpoint is saved. EVAL_PERIOD specifies the number of steps after which the eval is run.
• OUTPUT_DIR stores the checkpoints and the tensorboard logs. --resume resumes the training from the last checkpoint stored in OUTPUT_DIR. If no checkpoint is present, it loads the weights from MODEL.WEIGHTS

Model Zoo

ScanNet Instance Segmentation | Dataset | mAP | mAP@25 | Config | Checkpoint |:-:|:-:|:-:|:-:|:-:| | ScanNet val (ResNet50) | 47.8 | 83.6 | config | checkpoint | ScanNet val (Swin-B) | 50.0 | 83.6 | config | checkpoint | ScanNet test (Swin-B) | 47.7 | 86.2 | config | checkpoint

ScanNet Semantic Segmentation | Dataset | mIoU | Config | Checkpoint |:-:|:-:|:-:|:-:| | ScanNet val (ResNet50) | 73.3 | config | checkpoint | ScanNet val (Swin-B) | 77.8 | config | checkpoint | ScanNet test (Swin-B) | 74.4 | config | checkpoint

Joint 2D-3D on ScanNet and COCO | Model | mAP (ScanNet) | mAP25 (ScanNet) | mAP (COCO) | Config | Checkpoint |:-:|:-:|:-:|:-:|:-:|:-:| | ODIN | 49.1 | 83.1 | 41.2 | config | checkpoint

ScanNet200 Instance Segmentation | Dataset | mAP | mAP@25 | Config | Checkpoint |:-:|:-:|:-:|:-:|:-:| | ScanNet200 val (ResNet50) | 25.6 | 36.9 | config | checkpoint | ScanNet200 val (Swin-B) | 31.5 | 45.3 | config | checkpoint | ScanNet200 test (Swin-B) | 27.2 | 39.4 | config | checkpoint

ScanNet200 Semantic Segmentation | Dataset | mIoU | Config | Checkpoint |:-:|:-:|:-:|:-:| | ScanNet200 val (ResNet50) | 35.8 | config | checkpoint | ScanNet200 val (Swin-B) | 40.5 | config | checkpoint | ScanNet test (Swin-B) | 36.8 | config | checkpoint

AI2THOR Semantic and Instance Segmentation | Dataset | mAP | mAP@25 | mIoU | Config | Checkpoint |:-:|:-:|:-|:-:|:-:|:-:| | AI2THOR val (ResNet) | 63.8 | 80.2 | 71.5 | config | checkpoint | AI2RHOR val (Swin) | 64.3 | 78.6 | 71.4 | config | checkpoint

Matterport3D Instance Segmentation | Dataset | mAP | mAP@25 | Config | Checkpoint |:-:|:-:|:-|:-:|:-:| | Matterport3D val (ResNet) | 11.5 | 27.6 | config | checkpoint | Matterport val (Swin) | 14.5 | 36.8 | config | checkpoint

Matterport3D Semantic Segmentation | Dataset | mIoU | mAcc | Config | Checkpoint |:-:|:-:|:-:|:-:|:-:| | Matterport3D val (ResNet) | 22.4 | 28.5 | config | checkpoint | Matterport3D val (Swin) | 28.6 | 38.2 | config | checkpoint

S3DIS Instance Segmentation | Dataset | mAP | mAP@25 | Config | Checkpoint |:-:|:-:|:-:|:-:|:-:| | S3DIS Area5 (ResNet50-Scratch) | 36.3 | 61.2 | config | checkpoint | S3DIS Area5 (ResNet50-Fine-Tuned) | 44.7 | 67.5 | config | checkpoint | S3DIS Area5 (Swin-B) | 43.0 | 70.0 | config | checkpoint

S3DIS Semantic Segmentation | Dataset | mIoU | Config | Checkpoint |:-:|:-:|:-:|:-:| | S3DIS (ResNet50) | 59.7 | config | checkpoint | Swin-B | 68.6 | config | checkpoint

Training Logs:

Please find training logs for all models here

Citing ODIN

If you find ODIN useful in your research, please consider citing:

@inproceedings{jain2024odin,
  title={Odin: a single model for 2D and 3D segmentation},
  author={Jain, Ayush and Katara, Pushkal and Gkanatsios, Nikolaos and Harley, Adam W and Sarch, Gabriel and Aggarwal, Kriti and Chaudhary, Vishrav and Fragkiadaki, Katerina},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={3564--3574},
  year={2024}
}

License

Shield:

The majority of ODIN is licensed under a MIT License.

Acknowledgement

Parts of this code were based on the codebase of Mask2Former and Mask3D.