Skip to content

aby2s/sharpmask

Repository files navigation

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

Example