SAR-Net: Shape Alignment and Recovery Network for Category-level 6D Object Pose and Size Estimation

📝 Overview

This repository contains the PyTorch implementation of the paper "SAR-Net: Shape Alignment and Recovery Network for Category-level 6D Object Pose and Size Estimation" [PDF] [Supp] [arXiv]. Our approach could recover the 6-DoF pose and 3D size of category-level objects from the cropped depth image.

For more results and robotic demos, please refer to our Webpage.

💫 Follow-up works

Our follow-up ICCV2023 work:
- PourIt!🫗: Weakly-supervised Liquid Perception from a Single Image for Visual Closed-Loop Robotic Pouring [Project Page] [PDF] [arXiv] [Supp]

📚 Dependencies

Python >= 3.6
PyTorch >= 1.4.0
CUDA >= 10.1

⚙️ Installation

conda create -n sarnet python=3.6
conda activate sarnet

pip install -r requirements.txt

📊 Prepare Dataset

Download [camera_train_processed] that we have preprocessed.
Download [CAMERA/val], [Real/test], [gts], [obj_models] and [nocs_results] provided by NOCS.
Download [mrcnn_mask_results] provided by DualPoseNet.

Unzip and organize these files in ./data/NOCS and ./results/NOCS as follows:

data
└── NOCS
    ├── camera_train_processed
    ├── template_FPS
    ├── CAMERA
    │   ├── val 
    │   └── val_list.txt
    ├── Real
    │   ├── test 
    │   └── test_list.txt
    ├── gts
    │   ├── cam_val
    │   └── real_test
    └── obj_models
        ├── val
        └── real_test
   

results
└── NOCS
    ├── mrcnn_mask_results
    │   ├── cam_val
    │   └── real_test
    └── nocs_results
        ├── val
        └── real_test

python preprocess/shape_data.py
python preprocess/pose_data.py
python generate_json.py

🌟 NOTE

NOTE that there is a small bug in the original evaluation code of NOCS w.r.t. IOU. We fixed this bug in our evaluation code and re-evaluated our method. Also thanks Peng et al. for further confirming this bug.

🪄 Demo

python demo.py --config ./config_evaluate/nocs_real_mrcnn_mask.txt

The predict results will be saved in the folder of ./examples

⏳ Training

modified the ${gpu_id} in config_sarnet.py

# using a single GPU
e.g. gpu_id = '0' 

# using multiple GPUs
e.g. gpu_id = '0,1,2,3'

python train_sarnet.py

🪄 Evaluation

We provide the pretrain models of 3D-GCN(./checkpoints/NOCS/pretrained/seg3d) and SAR-Net(./checkpoints/NOCS/pretrained/sarnet) for fast evaluation.

python evaluate.py --config ./config_evaluate/nocs_real_mrcnn_mask.txt

We also provide the results reported in our paper for comparison.

🔖 Citation

If you find our work helpful, please consider citing:

@InProceedings{Lin_2022_CVPR,
    author    = {Lin, Haitao and Liu, Zichang and Cheang, Chilam and Fu, Yanwei and Guo, Guodong and Xue, Xiangyang},
    title     = {SAR-Net: Shape Alignment and Recovery Network for Category-Level 6D Object Pose and Size Estimation},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2022},
    pages     = {6707-6717}
}

🌹 Acknowledgment

Our implementation leverages the code from NOCS, SPD and 3DGCN. Thanks for the authors' work.

hetolin / SAR-Net