Skip to content

Latest commit

 

History

History
45 lines (34 loc) · 1.8 KB

BALLOON.md

File metadata and controls

45 lines (34 loc) · 1.8 KB

Balloon Demo

This is a demo on how to train tensorpack's Mask R-CNN on a custom dataset. We use the balloon dataset as an example.

  1. Download and unzip the dataset:
wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
unzip balloon_dataset.zip

  1. (included already) Since this dataset is not in COCO format, we add a new file dataset/balloon.py to load the dataset. Refer to dataset/dataset.py on the required interface of a new dataset.

  2. (included already) Register the names of the new dataset in train.py and predict.py, by calling register_balloon("/path/to/balloon_dataset")

  3. Download a model pretrained on COCO from tensorpack model zoo:

wget http://models.tensorpack.com/FasterRCNN/COCO-MaskRCNN-R50FPN2x.npz
  1. Start fine-tuning on the new dataset:
./train.py --config DATA.BASEDIR=~/data/balloon MODE_FPN=True \
	"DATA.VAL=('balloon_val',)"  "DATA.TRAIN=('balloon_train',)" \
	TRAIN.BASE_LR=1e-3 TRAIN.EVAL_PERIOD=0 "TRAIN.LR_SCHEDULE=[1000]" \
	"PREPROC.TRAIN_SHORT_EDGE_SIZE=[600,1200]" TRAIN.CHECKPOINT_PERIOD=1 DATA.NUM_WORKERS=1 \
	--load COCO-MaskRCNN-R50FPN2x.npz --logdir train_log/balloon
  1. You can train as long as you want, but it only takes a few minutes to produce nice results. You can visualize the results of the latest model by:
./predict.py --config DATA.BASEDIR=~/data/balloon MODE_FPN=True \
	"DATA.VAL=('balloon_val',)"  "DATA.TRAIN=('balloon_train',)" \
	--load train_log/balloon/checkpoint --predict ~/data/balloon/val/*.jpg

This command will produce images like this in your window:

demo