SurfEmb
SurfEmb: Dense and Continuous Correspondence Distributions
for Object Pose Estimation with Learnt Surface Embeddings
Rasmus Laurvig Haugard, Anders Glent Buch
IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2022
pre-print | project-site
The easiest way to explore correspondence distributions is through the project site.
The following describes how to reproduce the results.
Install
Download surfemb:
$ git clone https://github.com/rasmushaugaard/surfemb.git
$ cd surfemb
All following commands are expected to be run in the project root directory.
Install conda , create a new environment, surfemb, and activate it:
$ conda env create -f environment.yml
$ conda activate surfemb
Download BOP data
Download and extract datasets from the BOP site. Base archive, and object models are needed for both training and inference. For training, PBR-BlenderProc4BOP training images are needed as well, and for inference, the BOP'19/20 test images are needed.
Extract the datasets under data/bop
(or make a symbolic link).
Model
Download a trained model (see releases):
$ wget https://github.com/rasmushaugaard/surfemb/releases/download/v0.0.1/tless-2rs64lwh.compact.ckpt -P data/models
OR
Train a model:
$ python -m surfemb.scripts.train [dataset] --gpus [gpu ids]
For example, to train a model on T-LESS on cuda:0
$ python -m surfemb.scripts.train tless --gpus 0
Inference data
We use the detections from CosyPose's MaskRCNN models, and sample surface points evenly for inference.
For ease of use, this data can be downloaded and extracted as follows:
$ wget https://github.com/rasmushaugaard/surfemb/releases/download/v0.0.1/inference_data.zip
$ unzip inference_data.zip
OR
Extract detections and sample surface points
Surface samples
First, flip the normals of ITODD object 18, which is inside out.
Then remove invisible parts of the objects
$ python -m surfemb.scripts.misc.surface_samples_remesh_visible [dataset]
sample points evenly from the mesh surface
$ python -m surfemb.scripts.misc.surface_samples_sample_even [dataset]
and recover the normals for the sampled points.
$ python -m surfemb.scripts.misc.surface_samples_recover_normals [dataset]
Detection results
Download CosyPose in the same directory as SurfEmb was downloaded in, install CosyPose and follow their guide to download their BOP-trained detection results. Then:
$ python -m surfemb.scripts.misc.load_detection_results [dataset]
Inference inspection
To see pose estimation examples on the training images run
$ python -m surfemb.scripts.infer_debug [model_path] --device [device]
[device] could for example be cuda:0 or cpu.
Add --real
to use the test images with simulated crops based on the ground truth poses, or further add --detections
to use the CosyPose detections.
Inference for BOP evaluation
Inference is run on the (real) test images with CosyPose detections:
$ python -m surfemb.scripts.infer [model_path] --device [device]
Pose estimation results are saved to data/results
.
To obtain results with depth (requires running normal inference first), run
$ python -m surfemb.scripts.infer_refine_depth [model_path] --device [device]
The results can be formatted for BOP evaluation using
$ python -m surfemb.scripts.misc.format_results_for_eval [poses_path]
Either upload the formatted results to the BOP Challenge website or evaluate using the BOP toolkit.
Extra
Custom dataset: Format the dataset as a BOP dataset and put it in data/bop.