#change the number of classes in your dataset (I have 39)
#don't need to change anything else here
Change the fine_tune_checkpoint to where the model.cpkt file is stored. Typically this will be stored in your specific model file (mine is in object_detection/ssd_mobilenet_v2_coco_2018_03_29 because I am using the ssd mobilenet model).
#change this to be where your model.ckpt file is located
Update this section of code to how many steps/iterations you want your model to run. For a smaller more defined dataset, you may need less steps. For a larger, more diverse dataset, a larger number of steps may be better. The goal is to get the optimal number of steps that will generate the smallest loss.
#update the number of steps you'd like to train your model
Change the training input path to be where your train.record file is located. Also change the label map path to be where your labelmap.pbtxt file is located (in training folder).
#change this to where train.record is located
#change this to where labelmap.pbtxt is located
For the evaluation configuration, set the number of examples to how many test images you have.
#change this to the number of test images you have
Lastly, change the validation input path to be where val.record is located. Set the validation label map path to be where labelmap.pbtxt is located (same as for training).
Now that the configuration file is ready to go, it is time to train the model! This step will be time consuming depending on how many steps you choose. When you train your model, the model will save checkpoints to the training directory every once in a while. Once you are done training the model the training directory should look something like this:
This is sample output for the training directory after training the model. In this example, I trained with Faster rCNN, and only trained 2000 times.
We want to use the most recent/updated model. In this case it it is model.ckpt-2000 because I trained 2000 times.
Create Inference Graph
After training the model, we need to create an inference graph. This can be done (regardless of model chosen) with the
my_inference_graph.py file on my GitHub, and execute the code below (change the model.ckpt):
!python3 my_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path training/ssd_mobilenet_v2_coco.config \
--trained_checkpoint_prefix training/model.ckpt-2000 \ #change this to be the most recent model saved in training directory
Once you create the inference graph, your model is ready to be tested. Update the code below to reflect your dataset and where your files/test images are stored.
# Name of the directory containing the object detection module we're using
MODEL_NAME = 'inference_graph'
#test image name
IMAGE_NAME = 'stopyield.jpg'
# Grab path to current working directory
CWD_PATH = os.getcwd()
# Path to frozen detection graph .pb file, which contains the model that is used
# for object detection.
PATH_TO_CKPT = os.path.join(CWD_PATH ,MODEL_NAME,'frozen_inference_graph.pb')
# Path to label map file
PATH_TO_LABELS = os.path.join(CWD_PATH,'training','labelmap.pbtxt')
# Path to image
#my image is stored in /content/, which is why I used that instead of CWD_PATH
PATH_TO_IMAGE = os.path.join('/content/', IMAGE_NAME)
# Number of classes the object detector can identify
NUM_CLASSES = 39
After making these changes and running the remaining code, you should successfully be able to detect objects in your test image!
Output image after training object detection model
Feel free to run this code by downloading the files from my