HuaWeitong / REDE

Code for "REDE: End-to-end Object 6D Pose Robust Estimation Using Differentiable Outliers Elimination"

REDE

Overview

This is an implementation of the paper "REDE: End-to-end Object 6D Pose Robust Estimation Using Differentiable Outliers Elimination". [paper] [video]

Dependencies

• Python 3.8
• PyTorch 1.4
• CUDA 10.0

Installation

1. Set up python environment

pip install -r requirement.txt

2. Compile KNN

cd lib/KNN_CUDA
$ROOT=/path/to/REDE
export PATH=$ROOT/lib/KNN_CUDA:$PATH
make && make install

3. Compile SVD

cd lib/torch-batch-svd
python setup.py install

4. Compile FPS

cd lib/fps
python setup.py install

Datasets

Prepare the dataset

Download the preprocessed LineMOD dataset and supplement (including fuse data, keypoints and script for pose format conversion), put supplement into Linemod_preprocessed and run gt_to_npy.py.

Download the preprocessed Occlusion LineMOD dataset.

Download the YCB_Video Dataset.

Create the soft link

ln -s path/to/LineMOD datasets/linemod/data
ln -s path/to/Occlusion_LineMOD datasets/occlusion_linemod/data
ln -s path/to/YCB datasets/ycb/data

Compute FPS keypoints

python lib/fps/get_fps.py

Training

• LineMOD:

python tools/train_linemod.py --obj_id 1

• YCB:

python tools/train_ycb.py

Evaluation

You can download the trained checkpoints from Link (LineMOD, YCB) and put them to checkpoints.

• LineMOD:

python tools/eval_linemod.py

• Occlusion LineMOD:

python tools/eval_occ.py

• YCB:

python tools/eval_ycb.py