fidler-lab / polyrnn-pp-pytorch

PyTorch training/tool code for Polygon-RNN++ (CVPR 2018)
703 stars 106 forks source link
annotation cvpr2018 deep-learning instance-segmentation labelling polygon-rnn pytorch

PolygonRNN++

This is the official PyTorch reimplementation of Polygon-RNN++ (CVPR 2018). This repository allows you to train new Polygon-RNN++ models, and run our demo tool on local machines. For technical details, please refer to:

Efficient Interactive Annotation of Segmentation Datasets with Polygon-RNN++
David Acuna*, Huan Ling*, Amlan Kar*, Sanja Fidler (* denotes equal contribution)
CVPR 2018
[Paper] [Video] [Project Page] [Demo]

Where is the code?

To get the code, please signup here. We will be using GitHub to keep track of issues with the code and to update on availability of newer versions (also available on website and through e-mail to signed up users).

If you use this code, please cite:

@inproceedings{AcunaCVPR18,
title={Efficient Interactive Annotation of Segmentation Datasets with Polygon-RNN++},
author={David Acuna and Huan Ling and Amlan Kar and Sanja Fidler},
booktitle={CVPR},
year={2018}
}

@inproceedings{CastrejonCVPR17,
title = {Annotating Object Instances with a Polygon-RNN},
author = {Lluis Castrejon and Kaustav Kundu and Raquel Urtasun and Sanja Fidler},
booktitle = {CVPR},
year = {2017}
}

Contents

  1. Reproduction Results
  2. Environment Setup
  3. Tool
    1. Backend
    2. Frontend
  4. Testing Models
  5. Training Models
    1. Data
    2. Training MLE Model
    3. Training RL Model
    4. Training Evaluator
    5. Training GGNN

Results

These are the reproduction results from this repository as compared to the paper

Training Type Num first points LSTM Beam Size Before Now
MLE + Att 1 1 65.43 66.35
MLE + Att + RL 1 1 67.17 67.45
MLE + Att + Evaluator 5 1 69.72 71.05
MLE + Att + Evaluator 5 8 70.21 70.91
MLE + Att + Evaluator + GGNN 5 8 71.38 72.05
MLE + Att + Evaluator + GGNN 5 1 - 72.08
MLE + Att + Evaluator + GGNN (Shared Encoder) 5 8 - 72.22
MLE + Att + Evaluator + GGNN (Shared Encoder) 5 1 - 72.33

Note: Benchmarked forward pass speed for the tool (with 5 first points and 1 beam size) is 0.3 seconds per interaction on a TitanXp

Note: Shared Encoder refers to sharing the Resnet between the graph network and the convLSTM network. In the original paper, the two networks were kept separate.

Environment Setup

All the code has been run and tested on Ubuntu 16.04, Python 2.7.12, Pytorch 0.4.0, CUDA 9.0, TITAN X/Xp and GTX 1080Ti GPUs

Tool

Backend

Frontend

Note: Replace SimpleHTTPServer with http.server if you are using python3 for the server

Note: You can setup your own image directory by editing Tool/frontend/static/js/polygon.js and passing that path to Tool/tool.py from the command line. This image directory MUST contain the pre-defined images that are defined in Tool/frontend/index.html

Testing Models

python Scripts/prediction/generate_annotation.py --exp <path_to_corresponding_experiment> --reload <path_to_checkpoint> --output_dir <path_to_store_predictions>

Training Models

Data

Cityscapes

Custom Dataset

To train on your custom datasets, you have one of two options:

Training

Training MLE model

Training RL model

Training Evaluator

Training GGNN