sahithchada / NOCS_PyTorch

NOCS reimplementation in PyTorch
MIT License
26 stars 4 forks source link

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

resss

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:

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

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:

Datasets

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:

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