Home

Awesome

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

<table> <tr> <td><img src="examples/0000_image.png" alt="Image 1"></td> <td><img src="examples/0000_depth_color.png" alt="Image 2"></td> <td><img src="examples/d435_0000_pred.png" alt="Image 3"></td> </tr> </table>

architecture

šŸ“ 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

šŸ“š Dependencies

āš™ļø Installation

conda create -n sarnet python=3.6
conda activate sarnet

pip install -r requirements.txt

šŸ“Š Prepare Dataset

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

<table> <tr> <td><img src="examples/0000_image.png" alt="Image 1"></td> <td><img src="examples/0000_depth_color.png" alt="Image 2"></td> <td><img src="examples/d435_0000_pred.png" alt="Image 3"></td> </tr> </table>

ā³ 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.