aby2s / sharpmask

TensorFlow implementation of DeepMask and SharpMask

SharpMask

DeepMask 1 and SharpMask 2 implementation on Python/TensorFlow.

Introduction

The repository contains an implementation of DeepMask and SharpMask models. DeepMask model predicts class agnostic object mask and object score, which is positive if an object is centered and fully contained in an image. SharpMask is an extension of DeepMask architecture, which uses a top-down refinement module to compute more precise object mask proposal.

The implementation is based on TensorFlow official ResNet-v2 3 model implementation and requires pre-trained ResNet weights.

ResNet model implementation is copied from the official TensorFlow repository.

Note, that variables in ResNet model checkpoint are explicitly placed on a GPU device. So they won't be restored on a CPU.

Quick start

1. Clone this repository

   git clone https://github.com/aby2s/sharpmask.git

2. Download COCO-2017 train and validation sets from here
3. Unpack images into $COCO_PATH/images and annotations into $COCO_PATH/annotations, where $COCO_PATH is any directory
4. Prepare tfrecord files by running

    python prepare_data.py --coco_path $COCO_PATH --train_path $TRAIN_PATH --validation_path $VAL_PATH

   where $TRAIN_PATH and $VAL_PATH are directories to store training and validation tfrecord files respectivly.

Now you are ready either to evaluate a pre-trained model or to train your one.

Training

1. Download ResNet pre-trained weights from here
2. Run to train DeepMask and SharpMask sequentially

   python run_model.py --model all --train --train_path $TRAIN_PATH --validation_path $VAL_PATH  --summary_path $SUMMARY_PATH --checkpoint_path $CKPT_PATH --resnet_ckpt $RESNET_CKPT_PATH

3. To train only DeepMask run

   python run_model.py --model deepmask --train --train_path $TRAIN_PATH --validation_path $VAL_PATH  --summary_path $SUMMARY_PATH --checkpoint_path $CKPT_PATH --resnet_ckpt $RESNET_CKPT_PATH

4. To continue training or to train sharpmask from pre-trained deepmask use --restore flag

   python run_model.py --model sharpmask --restore --train --train_path $TRAIN_PATH --validation_path $VAL_PATH  --summary_path $SUMMARY_PATH --checkpoint_path $CKPT_PATH

Evaluation

1. To run evaluation computing mean IOU on compiled coco validation set:

   python run_model.py --model sharpmask --restore --validate --validation_path $VAL_PATH  --summary_path $SUMMARY_PATH --checkpoint_path $CKPT_PATH

2. To apply object mask and calculate score on a single image run

   python run_model.py --model sharpmask --restore --evaluate --eval_source $EVAL_SOURCE --eval_target $EVAL_TARGET  --summary_path $SUMMARY_PATH --checkpoint_path $CKPT_PATH

Pre-trained weights

Pre-trained weights are available here

Example