Code for "EIF-6D: Dual-Stream Transformation Network with Implicit-Explicit Fusion for Category-Level 6D Pose Estimation"

EIF6D introduces a novel approach to category-level 6D pose estimation,a dual-stream transformation network that integrates explicit and implicit collaborative deformation mechanisms, which achieves high-precision alignment between shape priors and observed objects at both point cloud level and feature level.

We provide a comprehensive conda environment configuration to ensure all dependencies are properly installed:

# Create and activate the conda environment
conda env create -f environment.yaml
conda activate eif6d

The environment includes:

• PyTorch 1.10.1 with CUDA 11.3 support
• Python 3.6.13

Compile the custom PointNet++ operations:

# Compile pointnet2 modules
cd model/pointnet2
python setup.py install

For the REAL275 and CAMERA25 datasets, please follow:

1. DPDN dataset instructions
2. SPD dataset instructions

The datasets should be organized in the following structure:

data
├── CAMERA
│   ├── train
│   └── val
├── camera_full_depths
│   ├── train
│   └── val
├── Real
│   ├── train
│   └── test
├── gts
│   ├── val
│   └── real_test
├── obj_models
│   ├── train
│   ├── val
│   ├── real_train
│   └── real_test
├── segmentation_results
│   ├── train_trainedwoMask
│   ├── test_trainedwoMask
│   └── test_trainedwithMask
└── mean_shapes.npy

Train the model from scratch using:

# For single GPU training
python trainPT2Net.py --gpus 0 --config config/PT2Net.yaml

# For multi-GPU training
python trainPT2Net.py --gpus 0,1 --config config/PT2Net.yaml

Training configurations can be adjusted in config/PT2Net.yaml.

Evaluate the trained model on test datasets:

python test_PT2.py --config config/PT2Net.yaml

Our method achieves state-of-the-art performance on the REAL275 and CAMERA25 benchmarks:

• Our code is developed upon DPDN and IST-Net.
• The dataset is provided by NOCS.