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>š 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.