Home

Awesome

SSD: Single Shot MultiBox Object Detector

SSD is an unified framework for object detection with a single network.

You can use the code to train/evaluate/test for object detection task.

Disclaimer

This is a re-implementation of original SSD which is based on caffe. The official repository is available here. The arXiv paper is available here.

This example is intended for reproducing the nice detector while fully utilize the remarkable traits of MXNet.

Due to the permission issue, this example is maintained in this repository separately. You can use the link regarding specific per example issues.

What's new

Demo results

demo1 demo2 demo3

mAP

ModelTraining dataTest datamAPNote
VGG16_reduced 300x300VOC07+12 trainvalVOC07 test77.8fast
VGG16_reduced 512x512VOC07+12 trainvalVOC07 test79.9slow
Inception-v3 512x512VOC07+12 trainvalVOC07 test78.9fastest
Resnet-50 512x512VOC07+12 trainvalVOC07 test78.9fast

Speed

ModelGPUCUDNNBatch-sizeFPS*
VGG16_reduced 300x300TITAN X(Maxwell)v5.11695
VGG16_reduced 300x300TITAN X(Maxwell)v5.1895
VGG16_reduced 300x300TITAN X(Maxwell)v5.1164
VGG16_reduced 300x300TITAN X(Maxwell)N/A836
VGG16_reduced 300x300TITAN X(Maxwell)N/A128
Forward time only, data loading and drawing excluded.

Getting started

sudo apt-get install python-opencv python-matplotlib python-numpy
# for Ubuntu/Debian
cp make/config.mk ./config.mk
# enable cuda, cudnn if applicable

Remember to enable CUDA if you want to be able to train, since CPU training is insanely slow. Using CUDNN is optional, but highly recommended.

Try the demo

# cd /path/to/mxnet-ssd
python demo.py --gpu 0
# play with examples:
python demo.py --epoch 0 --images ./data/demo/dog.jpg --thresh 0.5
python demo.py --cpu --network resnet50 --data-shape 512
# wait for library to load for the first time

Train the model

This example only covers training on Pascal VOC dataset. Other datasets should be easily supported by adding subclass derived from class Imdb in dataset/imdb.py. See example of dataset/pascal_voc.py for details.

cd /path/to/where_you_store_datasets/
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar
# Extract the data.
tar -xvf VOCtrainval_11-May-2012.tar
tar -xvf VOCtrainval_06-Nov-2007.tar
tar -xvf VOCtest_06-Nov-2007.tar
ln -s /path/to/VOCdevkit /path/to/mxnet/example/ssd/data/VOCdevkit

Use hard link instead of copy could save us a bit disk space.

# cd /path/to/mxnet/example/ssd
bash tools/prepare_pascal.sh
# or if you are using windows
python tools/prepare_dataset.py --dataset pascal --year 2007,2012 --set trainval --target ./data/train.lst
python tools/prepare_dataset.py --dataset pascal --year 2007 --set test --target ./data/val.lst --shuffle False
# cd /path/to/mxnet/example/ssd
python train.py
# note that a perfect training parameter set is yet to be discovered for multi-GPUs
python train.py --gpus 0,1,2,3 --batch-size 32

Evalute trained model

Make sure you have val.rec as validation dataset. It's the same one as used in training. Use:

# cd /path/to/mxnet/example/ssd
python evaluate.py --gpus 0,1 --batch-size 128 --epoch 0

Convert model to deploy mode

This simply removes all loss layers, and attach a layer for merging results and non-maximum suppression. Useful when loading python symbol is not available.

# cd /path/to/mxnet/example/ssd
python deploy.py --num-class 20

Convert caffe model

Converter from caffe is available at /path/to/mxnet/example/ssd/tools/caffe_converter

This is specifically modified to handle custom layer in caffe-ssd. Usage:

cd /path/to/mxnet/example/ssd/tools/caffe_converter
make
python convert_model.py deploy.prototxt name_of_pretrained_caffe_model.caffemodel ssd_converted
# you will use this model in deploy mode without loading from python symbol(layer names inconsistent)
python demo.py --prefix ssd_converted --epoch 1 --deploy

There is no guarantee that conversion will always work, but at least it's good for now.

Legacy models

Since the new interface for composing network is introduced, the old models have inconsistent names for weights. You can still load the previous model by rename the symbol to legacy_xxx.py and call with python train/demo.py --network legacy_xxx For example:

python demo.py --network 'legacy_vgg16_ssd_300.py' --prefix model/ssd_300 --epoch 0