Fransiska
Fransiska

Reputation: 43

Is there any way to modify available TensorFlow models architecture (such as ssd or fast r-cnn) so it is optimized for only one object detection?

I am new to Machine Learning and TensorFlow so I'm sorry and please correct me if my understanding is wrong. I have this project, developing a real time traffic-light detection with TensorFlow.

I've been working with pre-trained TensorFlow models such as SSD Mobilenet and Faster R-CNN Resnet. However, the expected accuracy result have not yet reached. I already considered to add some more data to the dataset (my dataset contains +/-1000 images), but because it is more work to add more data (since I have to do another data taking and label all images), which could take days. I want to consider another option.

Is there any way to modify TensorFlow models architecture so I could optimize and make it focused for only traffic light detection? I've been looking through TensorFlow models folder and could not find in which file these model architectures defined.

Any help will be appreciated. Thank you

Upvotes: 0

Views: 599

Answers (1)

Danny Fang
Danny Fang

Reputation: 4071

What you need is fine-tuning a pretrained model on your traffic light dataset. Given you already have about 1000 images for just one class, this is a descent dataset.

In order to perform fine-tuning, there are some important steps to do.

  1. First you need to transform your data into tfrecord format. Follow this tutorial to generate tfrecord files. This is actually a difficult step.
  2. Create a label_map.pbtxt for your model, since it is only traffic light, what you need is this. Here are sample label_map files.

    item { name: "traffic-light" id: 1 display_name: "traffic-light" }

  3. Then you need to prepare a pipeline config file for the model. Since you want to have a real-time detector, I suggest you use SSD-mobilenet models. Some sample config files are available here. You can take one of this sample configs and modify some fields to get the pipeline config for your model.

Suppose you choose ssd_mobilenet, then you can modify this config file, ssd_mobilenet_v2_coco.config. Specifically you need to modify these fields:

  • num_classes: you need to change from 90 to 1 since you only have traffic lights to detect.
  • input_path: (in both train_input_reader and eval_input_reader), point this to the tfrecord file you created.
  • label_map_path: (in both train_input_reader and eval_input_reader), point this to the label_map file you created.
  • fine_tune_checkpoint: set this path to a downloaded pretrained ssd-mobilenet model.

Depending on your training results, you may need to further adjust some of the fields in the config file, but after the training your model will focus only on traffic light class and will likely have a high accuracy.

All the specific tutorials can be found on the repo site. If you have any further questions, you can ask on stackoverflow with tag: object-detection-api and a lot people would help.

Upvotes: 0

Related Questions