In this work, we propose a new method for unseen object instance segmentation by learning RGB-D feature embeddings from synthetic data. A metric learning loss functionis utilized to learn to produce pixel-wise feature embeddings such that pixels from the same object are close to each other and pixels from different objects are separated in the embedding space. With the learned feature embeddings, a mean shift clustering algorithm can be applied to discover and segment unseen objects. We further improve the segmentation accuracy with a new two-stage clustering algorithm. Our method demonstrates that non-photorealistic synthetic RGB and depth images can be used to learn feature embeddings that transfer well to real-world images for unseen object instance segmentation. arXiv, Talk video

Unseen Object Clustering is released under the NVIDIA Source Code License (refer to the LICENSE file for details).

If you find Unseen Object Clustering useful in your research, please consider citing:

@inproceedings{xiang2020learning,
    Author = {Yu Xiang and Christopher Xie and Arsalan Mousavian and Dieter Fox},
    Title = {Learning RGB-D Feature Embeddings for Unseen Object Instance Segmentation},
    booktitle = {Conference on Robot Learning (CoRL)},
    Year = {2020}
}

Successfully tested with

• Ubuntu 20.04
• PyTorch 1.7.1+cu101
• CUDA 10.1

1. Install PyTorch.

2. Install python packages

   pip install -r requirement.txt

3. Install Python-PCL In order to use the python-pcl functionality to load datasets, you can not use the deprecated and un-maintained pip-package anymore. Instead install the apt-package python3-pcl with

   sudo apt install python3-pcl

   This will install the required PCL-Package in /usr/lib/python3/dist-packages/pcl/. You can copy this folder into the lib/python3.X/site-packages/-folder of any local virtual environment to make it accessible there.

   sudo cp -R /usr/lib/python3/dist-packages/pcl <path-to-virtualenv>/lib/python3.X/site-packages/

   If you get problems with access-rights of the package use this:

   sudo chown -R $USER <path-to-virtualenv>/lib/python3.X/site-packages/pcl/

• Download our trained checkpoints from here, save to $ROOT/data.

1. Download our trained checkpoints first.

2. Run the following script for testing on images under $ROOT/data/demo.

   ./experiments/scripts/demo_rgbd_add.sh

1. Download the Tabletop Object Dataset (TOD) from here (34G).

2. Create a symlink for the TOD dataset

   cd $ROOT/data
   ln -s $TOD_DATA tabletop

3. Training and testing on the TOD dataset

   cd $ROOT
   
   # multi-gpu training, we used 4 GPUs
   ./experiments/scripts/seg_resnet34_8s_embedding_cosine_rgbd_add_train_tabletop.sh
   
   # testing, $GPU_ID can be 0, 1, etc.
   ./experiments/scripts/seg_resnet34_8s_embedding_cosine_rgbd_add_test_tabletop.sh $GPU_ID $EPOCH
   

1. Download the OCID dataset from here, and create a symbol link:

   cd $ROOT/data
   ln -s $OCID_dataset OCID

2. Download the OSD dataset from here, and create a symbol link:

   cd $ROOT/data
   ln -s $OSD_dataset OSD

3. Check scripts in experiments/scripts with name test_ocid or test_ocd. Make sure the path of the trained checkpoints exist.

   experiments/scripts/seg_resnet34_8s_embedding_cosine_rgbd_add_test_ocid.sh
   experiments/scripts/seg_resnet34_8s_embedding_cosine_rgbd_add_test_osd.sh
   

• Python2 is needed for ROS.

• Make sure our pretrained checkpoints are downloaded.

  # start realsense
  roslaunch realsense2_camera rs_aligned_depth.launch tf_prefix:=measured/camera
  
  # start rviz
  rosrun rviz rviz -d ./ros/segmentation.rviz
  
  # run segmentation, $GPU_ID can be 0, 1, etc.
  ./experiments/scripts/ros_seg_rgbd_add_test_segmentation_realsense.sh $GPU_ID

Our example: