Use TensorFlow object detection API and MobileNet SSDLite model to train a pedestrian detector by using VOC 2007 + 2012 dataset
Forked from https://github.com/cftang0827/pedestrian_detection_ssdlite with many thanks :) This fork pins tensorfow version and includes dependencies.
Use pip install the package
pip install git+https://github.com/garyburgmann/pedestrian_detection_ssdlite
After installation, you can use the API:
import cv2 from pedestrian_detection_ssdlite import api from matplotlib import pyplot as plt img = cv2.imread('test_img/example.jpg') bbox_list = api.get_person_bbox(img, thr=0.6) print(bbox_list) for i in bbox_list: cv2.rectangle(img, i, i, (125, 255, 51), thickness=2) plt.imshow(img[:, :, ::-1]) plt.show()
and you will get the list of person bbox:
[[(267, 62), (343, 270)], [(201, 65), (255, 227)], [(187, 64), (228, 169)]]
Use the api and pretrained model that I provided
test.py, and I also provided a simple interface for using model, if you don't want to know the detail, please just copy whole api directory to your project and follow the way in
test.py, you will know how to use it.
The model was trained by using Tensorflow object detection API, and there are tons of good tutorial, so I would not explain too much here.
There are some useful tutorials link:
There are some classical scientific paper about modern object detection and the algorithm in Tensorflow Object Detection API
- RCNN: https://arxiv.org/pdf/1311.2524.pdf
- FastRCNN: https://www.cv-foundation.org/openaccess/content_iccv_2015/papers/Girshick_Fast_R-CNN_ICCV_2015_paper.pdf
- FasterRCNN: https://arxiv.org/pdf/1506.01497.pdf
- YOLO: https://pjreddie.com/media/files/papers/yolo.pdf
- SSD: https://www.cs.unc.edu/~wliu/papers/ssd.pdf
- MaskRCNN: https://research.fb.com/wp-content/uploads/2017/08/maskrcnn.pdf
git clone https://github.com/tensorflow/models.git
- Prepare VOC dataset (we took VOC 2012 + 2007 for example)
create_pascal_tf_record_only_person.pyand modify to the version that extract only one class, I used "person" here for example
pascal_label_map.pbtxtand put one class called "person"
- Make training and validation dataset to tfrecord format
- Find the proper model and detection algorithm for yourself, I took MobileNet SSDlite for example here. You can find some information in Tensorflow Object Detection API's website https://github.com/tensorflow/models/tree/master/research/object_detection
- Download proper pretrained model from tensorflow object detection API model zoo, because we usually DO NOT train the model from scratch, instead, we will use the pretrained model from Google and do transfer learning, you can find tons of pretrained model of Tensorflow Object Detection API here: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md
- Find the proper config file for your model and algorithm, put it into directory, I used
ssdlite_mobilenet_v2_coco.configfor example. You can find the config in
models/research/object_detection/samples/configs/, and modify some part (training dataset path, pretrained model ckpt path) of config file to your custom dataset.
./train.shand start train the model
- After training model, you will find there are many ckpt model in
models/research/object_detection/export_inference_graph.pyto freeze ckpt model to pb model
The tutorial is really rough, please take a look from others tutorial that I provided.