Custom Model Object Detection with OpenCV and ImageAI

• Download source - 6.5 KB

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()
detector.setModelTypeAsYOLOv3()
detector.setModelPath("hardhat\models\detection_model-ex-020--loss-0008.462.h5")
detector.setJsonPath("hardhat\json\detection_config.json")
detector.loadModel()

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)
showImage(convertedImage)

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

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()
person_detector.setModelTypeAsYOLOv3()
person_detector.setModelPath('yolo.h5')
person_detector.loadModel()

hard_hat_detector = CustomObjectDetection()
hard_hat_detector.setModelTypeAsYOLOv3()
hard_hat_detector.setModelPath("hardhat\models\detection_model-ex-020--loss-0008.462.h5")
hard_hat_detector.setJsonPath("hardhat\json\detection_config.json")
hard_hat_detector.loadModel()

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:

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!