Become a Marketing Hero

Become a Marketing Hero

Irresistible website designs and kickass marketing programs that deliver results

Learn More

Train Your Image Recognition AI With 5 Lines of Code

Computer vision has left the realm of science fiction and become science fact. No longer confined to the pages of academic journals and scholastic programming experiments, computer vision is quickly becoming mainstream in every industry and every walk of life. It is predicted that computer vision will be a  $48.6 billion industry by 2022. 

Image recognition is one of the most exciting innovations in the field of machine learning and artificial intelligence. Artificial intelligence is becoming increasingly essential for success in today's business world. Almost  two-thirds of CEOs believe artificial intelligence is helping them to pull ahead of their competition. 

With artificial intelligence becoming mainstream, this means that you no longer have to be an expert programmer or data scientist to deploy things like machine learning. With so many of the world's best developers working in the field, machine learning and computer vision are getting close to becoming a plug-and-play solution. 

We're going to show you how you can train image recognition AI with 5 lines of code, so you can implement it in your own business stack!

How To Train Your Image Recognition AI with 5 Lines of Code?

Computer vision is one of the most exciting and promising applications of machine learning and artificial intelligence. This is partly due to the fact that computer vision actually encompasses many different disciplines. Nearly all of them have profound  implications for businesses in a wide array of industries. 

Consider object detection, which is one of the first things to come to mind when someone asks "What is image recognition?". Object detection is widely used In Securit Self-Driving Vehicles, Pedestrian Detection, Industrial Automation and Digital. Marketing

The popularity of object detection means that there's an impressive array of resources that are available to help you train your own ML networks. There's the  Google Open Images dataset, for example. This gives you over 15 million images to use in your computer vision projects.

However, using public datasets comes with its own challenges. Training your own AI for image recognition still takes a bit of technical expertise. 

We're going to walk you through how to train your own image recognition AI with 5 lines of code. 

Step 1. Prepare Your Dataset

Custom image detection AI requires the right images for the job. These need to be put into a format your machine learning AI can understand. For this tutorial, we'll be using the  Pascal VOC format

Next, you'll have to decide what kind of objects you want to detect and you'll need to gather about 200 images of that object to train your image recognition AI. 

You'll then need to  install LabelImg to annotate the images. This allows you to save data from an image to an XML file. 

Now, create a directory for your images. Create a folder called 'images'. Then create another one which we’ll name 'annotations'. 

Now, run LabelImg and enable Auto Save Mode under the View menu. Select Open Dir from the top-left corner and then choose your images folder when prompted for a directory. 

You'll then need to set the directory where the image is saved. Select Change Save Dir in LabelImg and select your annotations folder. 

In LabelImg, you'll need to select the objects you're trying to detect. Click the 'Create RectBox' button on the bottom-left corner of the screen. Then draw a box around the object you want.

Enter the name of the object in the pop-up dialogue window and click the Okay button once you're done. 

You'll need to do this for all of the images in your images folder by selecting the 'Next Image' button and repeating the same process for the rest of the images in your images folder. Once you're done, your annotations folder will be full of XML files.

You'll need to do this twice. Once for the training model and once for the validation. Each one will have its own images and annotations folder. 

Step 2. Install ImageAI

We're using Python to implement our computer vision AI. We'll assume you've already got Python installed. If not, install the newest version of Python and then continue this tutorial. 

Now, you're going to install the libraries you'll need for your machine learning project. We're starting  with TensorFlow, which is one of the most popular Python libraries for machine learning. 

Run the following code in Terminal or the Command Prompt:

pip install tensorflow==2.4.0

Now, you'll install the other libraries you'll be using to train your image recognition AI. Run the following:

pip install keras==2.4.3 numpy==1.19.3 pillow==7.0.0 scipy==1.4.1 h5py==2.10.0 matplotlib==3.3.2 opencv-python keras-resnet==0.2.0

Finally, you're going to install a Python library called ImageAi.

pip install imageai --upgrade

We'll assume you've created a directory for this project in your programming environment. In that folder, create a Python file using your text editor of choice. Call it something like 

Copy this code into the Python file:

from imageai.Classification.Custom import ClassificationModelTrainer

model_trainer = ClassificationModelTrainer()



model_trainer.trainModel(num_objects=10, num_experiments=200, enhance_data=True, batch_size=32, show_network_summary=True)

Those 5 lines of code are all that you need to create your own image detection AI. 

However, let's take a closer look at the code to help you picture what's happening. 

Import ClassificationModelTrainer

Installing ImageAI makes a number of different modules available for your image recognition AI. The first line of code calls the ClassificationModelTrainer function. This makes it available to be used by the rest of the components. 


The next line of code creates a variable called Model_Trainer. The third states that it will use ResNes50 as the model type. Machine learning can use four different types for its models. They are:


  • MobileNetv2
  • ResNet50
  • InceptionV3
  • DenseNet151


Set Directory

The fourth line of code declares where the dataset directory lives. We've called ours 'objectdetect' but you can label your image recognition AI as you wish. 


The fifth line of code calls the trainModel function. It then defines a number of parameters.

number_objects declares how many objects are in your dataset. num_experiments determines how many times the model will be run. Enhance_data tells the machine learning AI if it needs to create enhanced copies of the original images to ensure accuracy. 

batch_size tells the machine learning model how many images to look at in one batch. show_network_summary creates a log of what your machine learning AI is doing. 

Now you're ready to run the program! Navigate to your programming folder and run:

$python Log

Since you've turned on show_network_summary, you'll see a log of what is doing. It should look something like this:


Total params: 23,608,202

Trainable params: 23,555,082

Non-trainable params: 53,120


Using Enhanced Data Generation

Found 4000 images belonging to 4 classes.

Found 800 images belonging to 4 classes.

JSON Mapping for the model classes saved to C:UsersUserPycharmProjectsFirstTrainingobjectdetectjsonmodel_class.json

Number of experiments (Epochs) : 200

Epoch 1/100

1/280 [>.............................] - ETA: 52s - loss: 2.3026 - acc: 0.2500

2/280 [>.............................] - ETA: 52s - loss: 2.3026 - acc: 0.2500

3/280 [>.............................] - ETA: 52s - loss: 2.3026 - acc: 0.2500




279/280 [===========================>..] - ETA: 1s - loss: 2.3097 - acc: 0.0625Epoch 00000: saving model to C:UsersUserPycharmProjectsFirstTrainingobjectdetectmodelsmodel_ex-000_acc-0.100000.h5


280/280 [==============================] - 51s - loss: 2.3095 - acc: 0.0600 - val_loss: 2.3026 - val_acc: 0.1000


The statement "JSON Mapping for the model classes saved to C:UsersUserPycharmProjectsFirstTrainingobjectdetectjsonmodel_class.json" shows that a file called model_class.json has been created for the objectdetect dataset. Once the code finishes running, you'll be able to use your dataset in other custom image recognition AI applications. 


The n/280 lines detail how many of the batches the machine learning AI has completed. 

Epoch 0000

The code `Epoch 00000: saving model to C:UsersUserPycharmProjectsFirstTrainingobjectmodelsmodel_ex-000_acc-0.100000.h5` refers to how many times the model has been run. The `acc-0.100000.h5` represents the accuracy of the model. 

This lets you know which model you should use for future predictions. The higher the accuracy, the more likely your AI is going to categorize images correctly. 

Create a Custom Image Recognition AI

You now have a fully-functional machine learning dataset. You can now use that dataset for future predictions. To test it out for yourself, create a new Python file in a new directory. We're calling ours

You'll also need to copy the JSON file created by your training model.

Copy this code into your new Python file: 

from imageai.Classification.Custom import CustomImageClassification

import os

execution_path = os.getcwd()

prediction = CustomImageClassification()





predictions, probabilities = prediction.predictImage("image.jpg", result_count=3)

for eachPrediction, eachProbability in zip(predictions, probabilities):

print(eachPrediction , " : " , eachProbability)

Run your new Python file:


Running this code will reveal the image classification and the probability of its accuracy. 

The ability to easily implement image recognition AI is a game changer for a number of different industries. Try it out for yourself and find out how technology can supercharge your business! 

Ready to Try Machine Learning for Yourself?

Image recognition will become more and more prevalent as the years progress. If you're not pulling ahead of your competition by now, you'll soon be falling behind... 

However, not everybody has the GPU resources to run a neural network at full speed. But don't worry - FluidStack is here to help! Please get in touch with one of our AI and machine learning specialists to help you find the right solution to your project.

 Our cloud-based GPUs are available for 3-5x cheaper than any other cloud provider! If you're ready to find out how image recognition AI can launch your business into the future,  contact one of our experts today!