Click here to Skip to main content
15,946,988 members
Articles / Artificial Intelligence / Machine Learning

Custom Model Object Detection with OpenCV and ImageAI

Rate me:
Please Sign up or sign in to vote.
5.00/5 (12 votes)
11 Jun 2020CPOL4 min read 26.3K   312   18   5
This article gives you a good starting point for your own object detection projects.
Here we'll look at using the trained model. We define our custom detection model, we create two detectors (one to detect all people, and one to detect people wearing hardhats), we compare the number of detections in each image to see if we caught someone not wearing a hardhat, and finally we briefly look at some other options ImageAI provides to improve the usability of 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 6 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 trained an ImageAI model to detect people wearing hardhats in an image.

Using the Trained Model

Now that we have a custom model ready to go, how can we actually use it? Let’s define some new code blocks similar to some of the ones we used previously. First, to define our custom detection model, create a code block like the following:

from imageai.Detection.Custom import CustomObjectDetection

detector = CustomObjectDetection()

This code block works in a similar way to our previous, pretrained model detector, except we define a model path and a JSON path for the configuration. This allows ImageAI to import the object names from the configuration file. Next, copy our previous code block that would detect people in images in our "people" directory and modify the code a little to detect them on our validation images:

import random
peopleImages = os.listdir("people")
randomFile = peopleImages[random.randint(0, len(peopleImages) - 1)]

detectedImage, detections = detector.detectObjectsFromImage(output_type="array", input_image="people/{0}".format(randomFile), minimum_percentage_probability=30)
convertedImage = cv.cvtColor(detectedImage, cv.COLOR_RGB2BGR)

for eachObject in detections:
    print(eachObject["name"] , " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"] )

Image 1

If you run these two code blocks now and look at the resulting image, you’ll see our model does a pretty reasonable job even with only 20 iterations. If you rerun the random image code block a few times, you’ll see that pretty much all the images with people in hardhats are picked up.

Next Steps

By taking advantage of two core libraries, OpenCV and ImageAI, we were able to use a pretrained object detection model, and to develop our own custom model, to detect if people are wearing hardhats.

We could combine these two models now and analyze images to ensure all the people within an image are wearing hardhats and, in a work setting, alert someone if they’re not. Here’s how to make that happen:

person_detector = ObjectDetection()

hard_hat_detector = CustomObjectDetection()

Here, we created two detectors: one to detect all people, and one to detect people wearing hardhats. Next, let’s load a random image and run it through both of our detectors:

peopleImages = os.listdir("people")
randomFile = peopleImages[random.randint(0, len(peopleImages) - 1)]

peopleOnly = person_detector.CustomObjects(person=True)

_, person_detections = detector.detectCustomObjectsFromImage(custom_objects=peopleOnly, output_type="array", input_image="people/{0}".format(randomFile), minimum_percentage_probability=30)

_, hard_hat_detections = detector.detectObjectsFromImage(output_type="array", input_image="people/{0}".format(randomFile), minimum_percentage_probability=30)

Finally, we’ll compare the number of detections in each image to see if we caught someone not wearing a hardhat:

if len(person_detections) > len(hard_hatdetections):
  print("hardhat non-compliance detected! Police have been dispatched!")
  # or perhaps just e-mail the boss to inform them of the violation

With that, we’ve successfully created a working hardhat detector.

There are a number of other things we could do to expand our fairly simple hardhat detection program. The next step would likely be performing more iterations to train our model. At the moment, the model does pick out people wearing hardhats pretty well, but the numbers mostly hover around the 50%-80% confidence level:

Image 2

By spending more time training the model, we could push these numbers higher, making our model more accurate.

If you have an NVIDIA GPU, you can dramatically speed up training by installing a GPU-enabled build of Tensorflow. If you don’t have a compatible GPU, you can access one inexpensively in the cloud. There are several good options:

  • Google Colab currently offers GPU-enabled notebook instances free of charge
  • Paperspace Gradient offers free GPU-enabled Jupyter Notebook instances
  • Azure Notebooks is another good choice, but you may not be able to use a GPU without paying unless you have Azure credits available

Some other options ImageAI provides could also improve the usability of the model:

  • The detectObjectsFromImage() method can take a parameter called extract_detected_objects=True. You could use this parameter to extract all the people from an image before iterating over them and testing to see if they’re wearing a hardhat. If they aren’t, their individual image could be sent to a manager.
  • ImageAI can also support video camera or live stream detection using the VideoObjectDetection class. This class works in a similar way to the ImageObjectDetection class and can be used to support near real-time video checking.

The ImageAI Documentation is pretty comprehensive and, hopefully, this article has given you a good starting point for your own object detection projects.

Wrapping Up

If you’ve been here since the first article in this series, it’s been quite a journey. But think of how far we’ve come! Starting from scratch, we built a hardhat detector that gets good results in the real world. And the best news is that the lessons we learned aren’t specific to hardhat detection. Using different sets of training data, you can train models to detect nearly anything!

This article is part of the series 'Workplace Safety with OpenCV View All


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

Written By
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

QuestionExcellent work Pin
fawadhitec28-Nov-20 7:27
fawadhitec28-Nov-20 7:27 
QuestionCustom trained model not detecting objects Pin
Member 1492322630-Aug-20 22:30
Member 1492322630-Aug-20 22:30 
AnswerRe: Custom trained model not detecting objects Pin
Glenn Prince31-Aug-20 7:58
professionalGlenn Prince31-Aug-20 7:58 
QuestionExcellent Series of Articles Pin
Member 1041561115-Jun-20 4:24
Member 1041561115-Jun-20 4:24 
QuestionGood work Pin
Xequence12-Jun-20 6:42
Xequence12-Jun-20 6:42 

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.