This repository contains the training code for the traffic lights detector used by the Kung-Fu-Panda team for the Udacity SDCND Capstone Project. The detector is based on the TensorFlow object detection API and uses the model produced by Alexey Simonov (our ex-team member) for classification.
You will need the TensorFlow models repository and a bunch of other things to run this code. Follow these installation instructions for a detailed explanation. Assuming you already have all the python dependencies installed, the set-up process may be summarised as follows:
sudo apt-get install protobuf-compiler
git clone [email protected]:tensorflow/models.git
cd models/research
protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
The create_sign_tf_record.py
program parses the input datasets and
produces input files for the object detection trainer. As far as the data
is concerned, I used the Bosch Small Traffic Lights and a bunch of
hand-annotated images extracted from the ROSBags provided by Udacity
and recorded from the simulator. I used this tool to create
annotations. The script rejects the images from the Bosch datasets containing
only very small objects.
cd data
unzip dataset_train_rgb.zip
tar zxf udacity-boxes.tar.gz
You'll need to convert the Bosch images to jpegs. I used ImageMagick and shell:
for i in *png; do convert $i `basename $i .png`.jpg; rm -f $i; done
First, you will need to download the base model:
mkdir model
cd model
wget http://download.tensorflow.org/models/object_detection/faster_rcnn_resnet101_coco_11_06_2017.tar.gz
tar zxf faster_rcnn_resnet101_coco_11_06_2017.tar.gz
cd ..
Then, you can train the model:
python $MODELS/research/object_detection/train.py \
--logtostderr \
--pipeline_config_path=faster_rcnn_resnet101_lights.config
--train_dir=train
I find that running it for about 30000 iterations produces pretty good results.
The most convenient way to use multiple models in a single application is to export them as static inference graphs. For the detection model, you can just use the utility provided by the TensorFlow Object Detection API:
python $MODELS/research/object_detection/export_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path faster_rcnn_resnet101_lights.config \
--trained_checkpoint_prefix train/model.ckpt-30452 \
--output_directory output_inference_graph.pb
I wrote a small utility program (export_classifier.py
) to export the
classification model created by Alexey.
I wrote a small testing script that takes a bunch of images as command-line
parameters and runs them, first, through the detector, and then, through
the classifier. ffmpeg
may then be used to create a video from multiple
frames:
./detect.py data/site/*.jpg
ffmpeg -framerate 24 -i output/left%04d.jpg output.mp4
You can watch sample videos by clicking on the images below: