This is a project of traffic lights detector and classification of there's red light visible.
The main purpose of this project was to address Red Litht Detection sub-project of the capstone project of Self-Driving Car Engineer Nanodegree by Udacity. The intent of this code is to be reaty to be used on CARLA autonomous car.
The traing code can be easily switched to detect other types of objects.
Feel free to experiment!
This is implementation in pytorch and there is no easy way to use this model for the Capstone project.
Keras implementation you might find useful is here:
github: bwosh/keras-traffic-lights
This is train and inference code of object detector for single class. It uses minimal version of implementation possible.
The code is ready to autimatically download COCO dataset annotations and extract photos for selected class.
Pre-trained model uses "Traffic light" class.
The code uses CenterNet implementation presented in "Objects as Points".
Used backbone model is DLA-34 proposed in "Deep Layer Aggregation".
- Implementation uses PyTorch 1.0+ with CUDA 8+
- Selected images from COCO dataset are used
- Object detector is CenterNet (much code taken from github.com/xingyizhou )
- Detector uses DLA-34 as backbone (minified version of code from github.com/xingyizhou
- PyTorch
- albumantations
- numpy
- opencv (for opening/saving images, drawing boxes)
cd src
pip install -r requirements.txt
cd src
python get_coco_images.py \
--class_name "traffic light" \
--output ../data/coco
cd src
python train.py \
--input ../data/coco \
--output ../models \
If you haven't trained model yourself you can download pretrained model and use for inference only:
- Download model from (512_map3950.pth)
- Place pth file in ./model folder
NOTICE: The dataset must be downloaded berore this step.
cd src
python train.py \
--input ../data/coco \
--output ../models \
--restore ../models/512_map3950.pth \
--val
cd src
python inference.py \
--folder [folder_with_images] \
--model ../models/512_map3950.pth
And enjoy the photos in the same folder but with '_output.jpg' suffix.