georgesung / ssd_tensorflow_traffic_sign_detection

Implementation of Single Shot MultiBox Detector in TensorFlow, to detect and classify traffic signs
MIT License
530 stars 222 forks source link
computer-vision object-detection python ssd tensorflow traffic-signs

SSD in TensorFlow: Traffic Sign Detection and Classification

Overview

Implementation of Single Shot MultiBox Detector (SSD) in TensorFlow, to detect and classify traffic signs. This implementation was able to achieve 40-45 fps on a GTX 1080 with an Intel Core i7-6700K.

Note this project is still work-in-progress. The main issue now is model overfitting. I am currently working on pre-training on VOC2012 first, then performing transfer learning over to traffic sign detection.

Currently only stop signs and pedestrian crossing signs are detected. Example detection images are below.

example1 example2 example3 example4 example5 example6

The model was trained on the LISA Traffic Sign Dataset, a dataset of US traffic signs.

Dependencies

How to run

Clone this repository somewhere, let's refer to it as $ROOT

Training the model from scratch:

Differences between original SSD implementation

Obivously, we are only detecting certain traffic signs in this implementation, whereas the original SSD implemetation detected a greater number of object classes in the PASCAL VOC and MS COCO datasets. Other notable differences are:

Performance

As mentioned above, this SSD implementation was able to achieve 40-45 fps on a GTX 1080 with an Intel Core i7 6700K.

The inference time is the sum of the neural network inference time, and Non-Maximum Suppression (NMS) time. Overall, the neural network inference time is significantly less than the NMS time, with the neural network inference time generally between 7-8 ms, whereas the NMS time is between 15-16 ms. The NMS algorithm implemented here has not been optimized, and runs on CPU only, so further effort to improve performance can be done there.

Dataset characteristics

The entire LISA Traffic Sign Dataset consists of 47 distinct traffic sign classes. Since we are only concered with a subset of those classes, we only use a subset of the LISA dataset. Also, we ignore all training samples where we do not find a matching default box, further reducing our dataset's size. Due to this process, we end up with very little data to work with.

In order to improve on this issue, we can perform image data augmentation, and/or pre-train the model on a larger dataset (e.g. VOC2012, ILSVRC)

Training process

Given the small size of our pruned dataset, I chose a train/validation split of 95/5. The model was trained with Adadelta optimizers, with the default parameters provided by TensorFlow. The model was trained over 200 epochs, with a batch size of 32.

Areas of improvement

There are multiple potential areas of improvement in this project: