Unseen Object Instance Segmentation for Robotic Environments

This is a PyTorch-based implementation of our network, UOIS-Net-3D, for unseen object instance segmentation. Our instance segmentation algorithm utilizes a two-stage method to explicitly leverage the strengths of depth and RGB separately for stronger instance segmentation. Surprisingly, our framework is able to learn from synthetic RGB-D data where the RGB is non-photorealistic. Details of the algorithm can be found in our arXiv paper:

Unseen Object Instance Segmentation for Robotic Environments
Christopher Xie, Yu Xiang, Arsalan Mousavian, Dieter Fox
IEEE Transactions on Robotics (T-RO), 2021.

Installation

We highly recommend setting up a virtual environment using Anaconda. Here is an example setup using these tools:

git clone https://github.com/chrisdxie/uois.git
cd uois3d/
conda env create -f env.yml

Models

You can find the models here. We provide a Depth Seeding Network (DSN) model trained on our synthetic Tabletop Object Dataset (TOD), a Region Refinement Network (RRN) model trained on TOD, and an RRN model trained on real data from the Google Open Images Dataset (OID).

Data

You can find the Tabletop Object Dataset (TOD) here. See the data loading and data augmentation code for more details.

Train the network

We provide sample training code in train_DSN.ipynb and train_RRN.ipynb.

Run the network

See uois_3D_example.ipynb for an example of how to run the network on example images. In order to run this file, Jupyter Notebook must be installed (this is included in env.yml). If you haven't used Jupyter Notebooks before, here is a tutorial to get you up to speed. This repository provides a few images in the example_images folder.

Notes:

Make sure to activate the Anaconda environment before running jupyter. This can be done with conda activate uois3d; jupyter notebook
the notebook should be run in the directory in which it lives (<ROOT_DIR>), otherwise the filepaths must be manually adjusted.
After downloading and unzipping the models, make sure to update checkpoint_dir in uois_3D_example.ipynb to point to the directory where the models live.

Citation

Our code is released under the MIT license.

If you find our work helpful in your research, please cite our work.

@article{xie2021unseen,
author    = {Christopher Xie and Yu Xiang and Arsalan Mousavian and Dieter Fox},
title     = {Unseen Object Instance Segmentation for Robotic Environments},
journal   = {IEEE Transactions on Robotics (T-RO)},
year      = {2021}
}

Scott-Su / uois