PyTorch Implementation of "Normalized Object Coordinate Space for Category-Level 6D Object Pose and Size Estimation"

Introduction

This is a PyTorch implementation of Normalized Object Coordinate Space for Category-Level 6D Object Pose and Size Estimation, a CVPR 2019 oral paper. The original Tensorflow implementation can be found here. The framework is built on a PyTorch implemenation of Mask-RCNN, which can be found here

For more details, please visit the project page

This repository includes:

Source code of NOCS.
Training code
Detection and evaluation code
Pre-trained weights

If you find this software helpful in your research or work, please consider citing using the Github cite tool on the right.

Contributions

This work was done as a final project for CSCI-5980: DeepRob taught by Professor Karthik Desingh. The contributors to the project are:

Special thanks to Bahaa Aldeeb for advising and troubleshooting help.

Requirements

This code has been tested with

CUDA 12.1 & 11.2
Python 3.10 & 3.9
PyTorch 2.0 & 1.10

Replicate the conda environment using:

conda create --name <env> --file requirements.txt

Implementation

This code implements the model of the original paper with the following settings:

NOCS values treated as classification (bins)
Unshared weights between NOCS heads
Symmetric Loss
Real & Synthetic data training (no COCO)

Datasets

CAMERA Dataset: Training/Test/IKEA_backgrounds/Composed_depths

+ Composed depth images contain the depths of both foreground synthetic objects and background real scenes for all train and validation data

Real Dataset: Training/Test
Ground truth pose annotation (for an easier evaluation): Val&Real_test
Object Meshes

You can download the files and store them under data/. The data folder general structure is shown:

.
└── data/
    ├── camera/
    │   ├── train
    │   └── val
    ├── real/
    │   ├── train
    │   └── test
    ├── obj_models/
    │   ├── real_test
    │   ├── real_train
    │   ├── train
    │   └── val
    ├── camera_full_depths/
    │   ├── train
    │   └── val
    ├── gts/
    │   ├── real_test
    │   └── val
    └── ikea_data

Pretrained weights

You can find the following checkpoints in this download link:

NOCS RCNN jointly trained on CAMERA, Real & MS COCO with 32 bin classification setting (Two sets of weights)
Mask RCNN pretrained on MS COCO dataset

You can download the checkpoints and store them under models/.

Training

# Train a new model from pretrained COCO weights
# Default uses same categories as original paper
python train.py

Detection Demo

# An image from synthetic validation or real test can be loaded and used for inference
# The image_id on line 89 can be changed to vary the image used.
python demo.py

Evaluation

# Weights performance on synthetic validation or real test can be evaluated
# The script must be run twice. First with detect set to False, and second with true (see line 87)
python demo_eval.py

sahithchada/NOCS_PyTorch