Click here to Skip to main content
15,030,448 members
Articles / Artificial Intelligence / Machine Learning
Posted 10 Jun 2020


6 bookmarked

Training a Custom Model with OpenCV and ImageAI

Rate me:
Please Sign up or sign in to vote.
5.00/5 (6 votes)
10 Jun 2020CPOL5 min read
In this article, we begin the process of creating a custom object detection model.
Here we look at some of the code behind training our model, and validating the model.

In this series, we’ll learn how to use Python, OpenCV (an open source computer vision library), and ImageAI (a deep learning library for vision) to train AI to detect whether workers are wearing hardhats. In the process, we’ll create an end-to-end solution you can use in real life—this isn’t just an academic exercise!

This is an important use case because many companies must ensure workers have the proper safety equipment. But what we’ll learn is useful beyond just detecting hardhats. By the end of the series, you’ll be able to use AI to detect nearly any kind of object in an image or video stream.

You’re currently on article 5 of 6:

  1. Installing OpenCV and ImageAI for Object Detection
  2. Finding Training Data for OpenCV and ImageAI Object Detection
  3. Using Pre-trained Models to Detect Objects With OpenCV and ImageAI
  4. Preparing Images for Object Detection With OpenCV and ImageAI
  5. Training a Custom Model With OpenCV and ImageAI
  6. Detecting Custom Model Objects with OpenCV and ImageAI

In the previous article, we cleaned our data and separated it into training and validation datasets.

Now we can begin the process of creating a custom object detection model. The general steps for training a custom detection model are:

  1. Train the model
  2. Validate the model; if validation is poor, tweak and retrain
  3. Visually test the model and the results
  4. Deploy the model

Training Our Model

Let's jump right into training our model. Create a new code block and enter the following:

from imageai.Detection.Custom import DetectionModelTrainer

trainer = DetectionModelTrainer()

trainer.setTrainConfig(object_names_array=["person hardhat"], batch_size=4, num_experiments=20, 


Image 1

This code block utilizes a new method in ImageAI's detection class, DetectionModelTrainer. The process for training any model is:

  • Define a new DetectionModelTrainer() method.
  • Set the model type as YOLOv3. At this point in time, YOLOv3 is the only model type that can be trained using ImageAI.
  • Set the directory that contains your data. Note that this folder must contain one folder named "train" and one named "validation." These folders must each contain one folder called "images" and another called "annotations."
  • Set the trainer’s configuration as follows:
    • Specify the names of the annotations used within the images. In our case, we use only "person hardhat."
    • Define the batch size, four in our case. This determines how many images the model will train in each batch. The higher your batch size, the better your model can be trained but, remember, the more powerful GPU you need.
    • Specify how many iterations of modelling to perform with num_experiments. The more iterations, the better your end result, but the longer it will take.
    • Specify, if desired, a pretrained model to transfer learning from, to give a better result quicker.
  • Start the model training process with trainModel().

The model will begin training and output a status for each cycle (or epoch). For each of these cycles, a loss is reported that determines if a model is better than in the previous cycle. If it is, that model will be saved, so make sure you have plenty of disk space free!

Validating the Model

Training a model can take a very long time. This model, with 20 iterations, took a bit over four hours to train. Some recommendations for training models suggest more than 200 hours. Before we leave our computer for several days to train a model, let's see what’s created after our model has been trained over 20 iterations.

Image 2

In your "hardhat" directory, you’ll see that some extra directories have been created: "cache," "json," "logs," and "models." The two important directories here are "json" and "models." The "json" directory contains the JSON configuration file required to use the model. The "model" directory contains a number of rather large model files with incremental numbers. Each of these files is the result of an iteration of model training that was better than the last.

So we have a number of models that are theoretically better and better, depending on the cycle number. Let’s test these out by validating them. Start a new code block and enter the following:

                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)

The only change you’ll need to make is the model path with the string: hardhat\models\detection_model-ex-020--loss-0008.462.h5, as each training run will be different. This method takes the following parameters:

  • model_path – specifies the model you wish to run a validation against
  • json_path – specifies the configuration file for the model training
  • iou_threshold – represents the ratio of intersection and union of the predicted and actual bounding box
  • object_threshold – the confidence level of detections to remove
  • nms_threshold – the confidence level when multiple bounding boxes are detected.

When we run this validation for a 20-iteration model, we get an average precision of 0.84464, or roughly 84%, which isn’t bad. But how does that compare against some of the others? Let's expand our code block to the following:

model05 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-005--loss-0014.238.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)
model10 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-010--loss-0011.053.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)
model15 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-015--loss-0009.620.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)
model20 = trainer.evaluateModel(model_path="hardhat\models\detection_model-ex-020--loss-0008.462.h5", 
                      json_path="hardhat\json\detection_config.json", iou_threshold=0.5, 
                      object_threshold=0.3, nms_threshold=0.5)

print('Iteration 05:', model05[0]['average_precision']['person hardhat'])
print('Iteration 10:', model10[0]['average_precision']['person hardhat'])
print('Iteration 15:', model15[0]['average_precision']['person hardhat'])
print('Iteration 20:', model20[0]['average_precision']['person hardhat'])

Image 3

This code block will take some time to run as it needs to load 4 different models, validate them, and save the results, so if you do run this, walk away and come back in a little while. When this code block eventually finishes, the last few lines will give you the results:

  • 5 Iterations – 71%
  • 10 Iterations – 78%
  • 15 Iterations – 83%
  • 20 Iterations – 84%

So, from these results, the more we iterate, the better our model gets. However, at some point, there are diminishing returns, so you need to take that into account when training your model as well.

Up Next

In this article, we learned how to train a custom image detection model using a prepared dataset.

Next, we’ll look at how to use this model to detect if people are wearing hardhats.


This article, along with any associated source code and files, is licensed under The Code Project Open License (CPOL)


About the Author

Glenn Prince
United States United States
Hi! I'm a Solution Architect, planning and designing systems based in Denver, Colorado. I also occasionally develop web applications and games, as well as write. My blog has articles, tutorials and general thoughts based on more than twenty years of misadventures in IT.

Comments and Discussions

QuestionOut Of Memory Error. Pin
Member 1502673721-Dec-20 2:03
MemberMember 1502673721-Dec-20 2:03 
QuestionImageai custom trained model not detecting Pin
Member 1492322613-Sep-20 5:49
MemberMember 1492322613-Sep-20 5:49 
QuestionImageai custom trained model not evaluating properly Pin
Member 1492322625-Aug-20 22:53
MemberMember 1492322625-Aug-20 22:53 
AnswerRe: Imageai custom trained model not evaluating properly Pin
Sean Ewington26-Aug-20 9:06
editorSean Ewington26-Aug-20 9:06 
According to the author:

Things to check would be:

1) The error is talking about no configuration data being found. When training the model a subdirectory in your data folder called "json" should be created with a "detection_config.json" file. Make sure that file exists and "json_path=" points to it

2) At least one "epoch" needs to complete, preferably more for an accurate model, to create the files you need.

3) Try deleting the cache," "json," "logs," and "models" folders and run the training again.

Sean Ewington

GeneralRe: Imageai custom trained model not evaluating properly Pin
Member 1492322627-Aug-20 21:22
MemberMember 1492322627-Aug-20 21:22 
GeneralRe: Imageai custom trained model not evaluating properly Pin
Member 1492322630-Aug-20 1:45
MemberMember 1492322630-Aug-20 1:45 
GeneralRe: Imageai custom trained model not evaluating properly Pin
Glenn Prince31-Aug-20 7:57
mvaGlenn Prince31-Aug-20 7:57 
GeneralRe: Imageai custom trained model not evaluating properly Pin
Member 149232263-Sep-20 21:30
MemberMember 149232263-Sep-20 21:30 
QuestionCustom Training Fails Due to Missing .jpg Files Pin
Member 1041561126-Jun-20 8:02
MemberMember 1041561126-Jun-20 8:02 
AnswerRe: Custom Training Fails Due to Missing .jpg Files Pin
Sean Ewington26-Jun-20 8:10
editorSean Ewington26-Jun-20 8:10 
GeneralRe: Custom Training Fails Due to Missing .jpg Files Pin
Member 1041561128-Jun-20 16:41
MemberMember 1041561128-Jun-20 16:41 
GeneralRe: Custom Training Fails Due to Missing .jpg Files Pin
Sean Ewington30-Jun-20 3:15
editorSean Ewington30-Jun-20 3:15 

General General    News News    Suggestion Suggestion    Question Question    Bug Bug    Answer Answer    Joke Joke    Praise Praise    Rant Rant    Admin Admin   

Use Ctrl+Left/Right to switch messages, Ctrl+Up/Down to switch threads, Ctrl+Shift+Left/Right to switch pages.