/PTSNet

Proposal, Tracking and Segmentation (PTS): A Cascaded Network for Video Object Segmentation

Primary LanguagePython

Proposal, Tracking and Segmentation (PTS): A Cascaded Network for Video Object Segmentation

By Qiang Zhou*, Zilong Huang*, Lichao Huang, Shen Han, Yongchao Gong, Chang Huang, Wenyu Liu, Xinggang Wang.(* means equal contribution)

This code is a implementation of the experiments on DAVIS dataset. For more detail, please refer to our paper.

Architecture


Overview of our proposed PTSNet for video object segmentation. OPN is designed for generating proposals of the interested objects and OTN aims to distinguish which one of the proposals is the best. Finally, DRSN does the final pixel level tracking(segmentation) task. Note in our implementation we couple OPN and OTN as a whole network, and spearate DRSN out under engineering consideration.

Usage

Preparation

  1. Install PyTorch 0.4.1 and necessary libraries like opencv, PIL etc. If there are many researchers being interested on PTSNet, we will re-implement the whole framework using PyTorch 1.0.1 or higher.

  2. There are some native CUDA implementations, InPlace-ABN and MaskRCNN Operators, which must be compiled at the very start.

    # Before you compile, you need to figure out several things:
# - The CUDA kernels supported by your GPU, here we use `sm_52`, `sm_61` and `sm_70` for NVIDIA Titan V.
# - `cuda` and `nvcc` paths in your operating system, which exist usually in `/usr/local/cuda` and `/usr/local/cuda/bin/nvcc` respectively.
# InPlace-ABN
cd model/inplace_ABN
bash build.sh
python3 build.py
# MaskRCNN Operators
cd coupled_otn_opn/tracking/maskrcnn/lib
bash make.sh

  3. You can train PTSNet from scratch or just evaluate our pretrained model.

    • Train it from scratch, you need to download:

       # DRSN: wget "https://download.pytorch.org/models/resnet50-19c8e357.pth" -O drsn/init_models/resnet50-19c8e357.pth
 # OPN: wget "https://s3-us-west-2.amazonaws.com/detectron/ImageNetPretrained/25093814/X-152-32x8d-IN5k.pkl" -O coupled_otn_opn/tracking/maskrcnn/data/X-152-32x8d-IN5k.pkl
 # If you want to use our pretrained OTN:
 #   wget https://drive.google.com/open?id=12bF1dRlEUZoQz3Qcr2WD3ojqNHzbCrjf, put it into `coupled_otn_opn/models/mdnet_davis_50cyche.pth`
 # Else please modify from py-MDNet(https://github.com/HyeonseobNam/py-MDNet) to train OTN on DAVIS by yourself.

    • If you want to use our pretrained model to do the evaluation, you need to download:

       # DRSN: https://drive.google.com/open?id=116yXnqX43BZ7kEgdzUhIeTSn1dbvcE2F, put it into `drsn/snapshots/drsn_yvos_10w_davis_3p5w.pth`
 # OPN: wget "https://s3-us-west-2.amazonaws.com/detectron/ImageNetPretrained/25093814/X-152-32x8d-IN5k.pkl" -O coupled_otn_opn/tracking/maskrcnn/data/X-152-32x8d-IN5k.pkl
 # OTN: https://drive.google.com/open?id=12bF1dRlEUZoQz3Qcr2WD3ojqNHzbCrjf, put it into `coupled_otn_opn/models/mdnet_davis_50cycle.pth`

  4. Dataset

    • YouTube-VOS: Download from YouTube-VOS, note we only need the training part(train_all_frames.zip), totally about 41G. Unzip, move and rename it to drsn/dataset/yvos.
    • DAVIS: Download from DAVIS, note we only need the 480p version(DAVIS-2017-trainval-480p.zip). Unzip, move and rename it to drsn/dataset/DAVIS/trainval and coupled_otn_opn/DAVIS/trainval. Here you need to make a subdirectory of trainval directory to store the dataset.

    And make sure to put the files as the following structure:

    .
├── drsn
│   ├── dataset
│   │   ├── DAVIS
│   │   │   └── trainval
│   │   │       ├── Annotations
│   │   │       ├── ImageSets
│   │   │       └── JPEGImages
│   │   └── yvos
│   │       └── train_all_frames
│   ├── init_model
│   │   └── resnet50-19c8e357.pth
│   └── snapshots
│       └── drsn_yvos_10w_davis_3p5w.pth
└── coupled_otn_opn
    ├── DAVIS
    │   └── trainval
    ├── models
    │   └── mdnet_davis_50cycle.pth
    └── tracking
        └── maskrcnn
            └── data
                └── X-152-32x8d-FPN-IN5k.pkl

Train and Evaluate

  • Firstly, check the directory of coupled_otn_opn and follow the README.md inside to generate our proposals. You can also skip this step for we have provided generated proposals in drsn/dataset/result_davis directory.
  • Secondly, enter drsn and check do_train_eval.sh to train and evaluate.
  • Finally, we also provide result masks by our PTSNet in result-masks-GoogleDrive. The quantitative results are measured by DAVIS official matlab toolbox.
J Mean F Mean G Mean
Avg 71.6 77.7 74.7

Acknowledgment

The work was mainly done during an internship at Horizon Robotics.

Thanks to the Third Party Libs