IMPORTANT NOTE: This project is for research and is not entirely stable, yet. A note will be posted here for when it is ready for more general use. Issues and/or Pull Requests are most welcome!

Notes

Quick Note

This PyTorch for SSD update is a staging environment for changes to the original work and, itself, a work in progress. It is a remake of https://github.com/amdegroot/ssd.pytorch and https://github.com/hli2020/object_detection.

Added to original repo:

• Custom dataset loader - custom.py (VGG Image Annotator input)
• Updates to working order of train.py and assoicated files
• Added the latest recommendation for specifying a GPU/CUDA device (.to(device) for nets and variables
• Templated structure for train and test data and dealing with more consistently in code

It has been tested with PyTorch 0.4.1 and a single class.

Quickstart

1. Install Python packages from requirements.txt
2. Label data with VGG Image Annotator using the rectangle option (we are doing object detection here to create a tiny SSD model for devices) and place in data/image_data/train (images and via .json file)
3. Start the Visdom server (to view results - test only right now) with: python -m visdom.server -port 8090
4. Run the training script, e.g., python train.py --experiment_name ssd_custom --dataset Custom --base_save_folder run --num_workers 0 --ssd_dim 300 --batch_size 4 --lr 1e-4 --max_epoch 50 --pretrain_model weights/vgg16_reducedfc.pth
   • The learning rate may need to be adjusted, larger or smaller, otherwise nan's are produced in losses.
   • More epochs may need to be added for better performance.
   • The learning rate scheduler may need to adjusted in the train.py file.
5. Test model with the test script (if retesting, may need to delete test folder, for detections to run again), e.g., python test.py --experiment_name ssd_custom --dataset Custom --base_save_folder run --num_workers 0 --ssd_dim 300 --trained_model run/ssd_custom/train/debug_ssd300_CUSTOM_epoch_49_iter_3.pth --prior_config custom --conf_thresh 0.05
6. Navigate to http://localhost:8090 to view results with the Visdom server.

Status

• Clean up files - remove any not needed
• Custom data loader for multiple annotators (as in custom.py)
  • VGG Image Annotator (http://www.robots.ox.ac.uk/~vgg/software/via/)
  • Reader for Open Images v4 boxes format
  • Reader for Visual Object Tagging Tool (https://github.com/Microsoft/VoTT)
• Fix the NaN loss issue for train.py on custom data (note: lowering the learning rate can elimate this as well)
• Remove the usage of args.cuda (args.cuda = not args.no_cuda and torch.cuda.is_available() or similar)
• Add better docstrings and comments to train.py and test.py
• Ensure still functioning on original COCO and VOC data as per original repo
• Update other files:
  • test.py
    • Get rid of imgsetpath = os.path.join(args.dataset_root, 'imagenames.txt') in favor of just glob
  • Working live.py code

For custom, make sure to update the number of iterations in the config.py (max_iter) (if minibatch size is 1, max_iter equals the number of input images in the training set).

Here is the original project's excellent Readme as of 2018-03-30 (for Installation, Models, etc.):

SSD: Single Shot MultiBox Object Detector, in PyTorch

A PyTorch implementation of Single Shot MultiBox Detector from the 2016 paper by Wei Liu, Dragomir Anguelov, Dumitru Erhan, Christian Szegedy, Scott Reed, Cheng-Yang, and Alexander C. Berg. The official and original Caffe code can be found here.

Table of Contents

• Installation
• Datasets
• Train
• Evaluate
• Performance
• Demos
• Future Work
• Reference

       

Installation

• Install PyTorch by selecting your environment on the website and running the appropriate command.
• Clone this repository.
  • Note: We currently only support Python 3+.
• Then download the dataset by following the instructions below.
• We now support Visdom for real-time loss visualization during training!
  • To use Visdom in the browser:

  # First install Python server and client
  pip install visdom
  # Start the server (probably in a screen or tmux)
  python -m visdom.server

  • Then (during training) navigate to http://localhost:8097/ (see the Train section below for training details).
• Note: For training, we currently support VOC and COCO, and aim to add ImageNet support soon.

Datasets

To make things easy, we provide bash scripts to handle the dataset downloads and setup for you. We also provide simple dataset loaders that inherit torch.utils.data.Dataset, making them fully compatible with the torchvision.datasets API.

COCO

Microsoft COCO: Common Objects in Context

Download COCO 2014

# specify a directory for dataset to be downloaded into, else default is ~/data/
sh data/scripts/COCO2014.sh

VOC Dataset

PASCAL VOC: Visual Object Classes

Download VOC2007 trainval & test

# specify a directory for dataset to be downloaded into, else default is ~/data/
sh data/scripts/VOC2007.sh # <directory>

Download VOC2012 trainval

# specify a directory for dataset to be downloaded into, else default is ~/data/
sh data/scripts/VOC2012.sh # <directory>

Training SSD

• First download the fc-reduced VGG-16 PyTorch base network weights at: https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth
• By default, we assume you have downloaded the file in the ssd.pytorch/weights dir:

mkdir weights
cd weights
wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

• To train SSD using the train script simply specify the parameters listed in train.py as a flag or manually change them.

python train.py

• Note:
  • For training, an NVIDIA GPU is strongly recommended for speed.
  • For instructions on Visdom usage/installation, see the Installation section.
  • You can pick-up training from a checkpoint by specifying the path as one of the training parameters (again, see train.py for options)

Evaluation

To evaluate a trained network:

python eval.py

You can specify the parameters listed in the eval.py file by flagging them or manually changing them.

Performance

VOC2007 Test

mAP

Original	Converted weiliu89 weights	From scratch w/o data aug	From scratch w/ data aug
77.2 %	77.26 %	58.12%	77.43 %

FPS

GTX 1060: ~45.45 FPS

Demos

Use a pre-trained SSD network for detection

Download a pre-trained network

• We are trying to provide PyTorch state_dicts (dict of weight tensors) of the latest SSD model definitions trained on different datasets.
• Currently, we provide the following PyTorch models:
  • SSD300 trained on VOC0712 (newest PyTorch weights)
    • https://s3.amazonaws.com/amdegroot-models/ssd300_mAP_77.43_v2.pth
  • SSD300 trained on VOC0712 (original Caffe weights)
    • https://s3.amazonaws.com/amdegroot-models/ssd_300_VOC0712.pth
• Our goal is to reproduce this table from the original paper

Try the demo notebook

• Make sure you have jupyter notebook installed.
• Two alternatives for installing jupyter notebook:

  1. If you installed PyTorch with conda (recommended), then you should already have it. (Just navigate to the ssd.pytorch cloned repo and run): jupyter notebook

  2. If using pip:

# make sure pip is upgraded
pip3 install --upgrade pip
# install jupyter notebook
pip install jupyter
# Run this inside ssd.pytorch
jupyter notebook

• Now navigate to demo/demo.ipynb at http://localhost:8888 (by default) and have at it!

Try the webcam demo

• Works on CPU (may have to tweak cv2.waitkey for optimal fps) or on an NVIDIA GPU
• This demo currently requires opencv2+ w/ python bindings and an onboard webcam
  • You can change the default webcam in demo/live.py
• Install the imutils package to leverage multi-threading on CPU:
  • pip install imutils
• Running python -m demo.live opens the webcam and begins detecting!

TODO

We have accumulated the following to-do list, which we hope to complete in the near future

• Still to come:
  • Support for the MS COCO dataset
  • Support for SSD512 training and testing
  • Support for training on custom datasets

Authors

• Max deGroot
• Ellis Brown

Note: Unfortunately, this is just a hobby of ours and not a full-time job, so we'll do our best to keep things up to date, but no guarantees. That being said, thanks to everyone for your continued help and feedback as it is really appreciated. We will try to address everything as soon as possible.

References

• Wei Liu, et al. "SSD: Single Shot MultiBox Detector." ECCV2016.
• Original Implementation (CAFFE)
• A huge thank you to Alex Koltun and his team at Webyclip for their help in finishing the data augmentation portion.
• A list of other great SSD ports that were sources of inspiration (especially the Chainer repo):
  • Chainer, Keras, MXNet, Tensorflow