Finding Training Data for OpenCV and ImageAI Object Detection

• 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 2 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 set up all of the development tools we’ll need to detect hardhat safety compliance with OpenCV.

But development tools aren’t enough. To train AI computer vision models, we also need a set of training data containing images of people, some with hardhats and some without. This data will let us teach our model to distinguish between people who are wearing hardhats and those who aren’t.

There are three general steps to create an object detection model:

1. Load samples of the data to classify
2. Train the model on the sample data
3. Test the model on different sample data containing both matches and non-matches

FindingTraining Images

We’re going to use an image database called ImageNet to get most of our source data for this stage. ImageNet maintains an image database, organized using a WordNet hierarchy of nouns. This database, for instance, allows us to retrieve all images in the hierarchy of:

artefact -> covering -> clothing -> headdress / headgear -> helmet -> hardhat

ImageNet also has a public API that returns a list of image URLs based on the noun ID. Let’s use this API to download lots of images of people in hardhats. The form of that API is:

http://www.image-net.org/api/text/imagenet.synset.geturls?wnid=<WordNet ID>

We’re going to be working with two WordNet nouns: the hardhat noun (as mentioned above) with ID n03492922; and the misc -> people noun with ID n07942152.

Downloading Hardhat Images

To download all the hardhat images from the API, create the following code block:

hardhatLoc = 'http://www.image-net.org/api/text/imagenet.synset.geturls?wnid=n03492922'

hardhatImages = req.get(hardhatLoc).text
noOfImages = 0

if not os.path.exists('hardhat'):
    os.makedirs('hardhat')

for i in hardhatImages.split('\n'):
    try:
        r = req.get(i, timeout=0.5)
        file = i.split("/")[-1].split('\r')[0]
        if 'image/jpeg' in r.headers['Content-Type']:
            if len(r.content) > 8192:
                with open('hardhat\\' + file, 'wb') as outfile:
                    outfile.write(r.content)
                    noOfImages += 1
                    print('Success: ' + file)
            else:
                print('Failed: ' + file + ' -- Image too small')
        else:
            print('Failed: ' + file + ' -- Not an image')

    except Exception as e:
        print('Failed: ' + file + ' -- Error')
        
print('*********** Download Finished **************')

There’s a lot going on here, so let’s step through it. First, we set the URL of the API we want to call using the variable hardhatLoc. Then, we query the API using the req variable that points to our requests library and we chain two methods together using this library:

• .get – takes a string URL and returns the response
• .text – takes a binary response and converts it into a string value

Once we have a list of source images from the API, we check to see if we have a folder to download the images into using the method os.path.exists. If we don't have a directory to put them in, the method .makedirs() will create the directory.

Next, we start our for loop, which uses the URLs we downloaded from the API. We split the string at every newline character (\n) and push that string into the variable i. To catch any errors downloading the file, we also open a try block.

Within the try block, this process will iterate over all the URLs we downloaded from the API. For each loop or URL, the following process occurs:

• We use .get() to try and download the file. We also use the optional parameter timeout to stop trying to download the file after 0.5 seconds
• We then split the URL string twice using .split().
• The first split occurs every "/" and the -1 value in the array gets the last value.
• The second split removes the trailing carriage return.
• The next step checks to see if the file downloaded is a jpeg using the .headers array. This array contains all the HTML headers from a response.
• After confirming it’s a jpeg, we then check to make sure the file is a decent size so we aren't processing any error images.
• If the downloaded file matches our requirements, we use the with statement block to save the file to the hardhat folder.
• Finally, we close all our if and try blocks with statements to print any errors.

Let's run this code block now. You’ll see that the code runs and prints out the success or failure of each image downloaded via the API. This may take a few minutes so let it run and wait for the complete message.

Checking the Downloaded Data

You’ll notice a few failure messages after running the downloader code block. This is good, as we want to have accurate images in our folder for the classifier to use. Open up an explorer window and take a look at the folder that contains your Jupyter notebook. You should now have a subfolder called "hardhat."

Open this hardhat folder and check the downloaded images. I ended up with 687 images of hardhats, mainly people wearing them but some with just the hardhat by itself. Check to ensure you have plenty of images.

Downloading People Images

Now that we have lots of images of hardhats, let’s copy our code block and download a sample set of people not wearing hardhats. Duplicate the existing code block and make the following changes:

• Change hardhat to people wherever it’s used, including variable names and text strings.
• Change the API id from n03492922 to n07942152.

Your code should now look like this:

peopleLoc = 'http://www.image-net.org/api/text/imagenet.synset.geturls?wnid=n07942152'

peopleImages = req.get(peopleLoc).text
noOfImages = 0

if not os.path.exists('people'):
    os.makedirs('people')

for i in peopleImages.split('\n'):
    try:
        r = req.get(i, timeout=0.5)
        file = i.split("/")[-1].split('\r')[0]
        if 'image/jpeg' in r.headers['Content-Type']:
            if len(r.content) > 8192:
                with open('people\\' + file, 'wb') as outfile:
                    outfile.write(r.content)
                    noOfImages += 1
                    print('Success: ' + file)
            else:
                print('Failed: ' + file + ' -- Image too small')
        else:
            print('Failed: ' + file + ' -- Not an image')

    except Exception as e:
        print('Failed: ' + file + ' -- Error')
        
print('*********** Download Finished **************')

When we run this code block, it should work just as it did before. But this time, if we check the folder where our Jupyter notebook is saved, there should be the "people" folder containing general images of groups of people with no hardhats on.

Up Next

We now have a dataset we can use to train and test AI models for hardhat detection.

Next, we’ll take a look at the results we can get using our images with a pretrained model.