How To Create Your Own Custom Image Classifier And Deploy It To Production

Source: Deep Learning on Medium

How To Create Your Own Custom Image Classifier And Deploy It To Production

In this article, we are going to make our own custom image classifier. We are going to use the library to train our deep learning model. It is a high-level library build on PyTorch, which allows us to build models using only a few lines of code.

And we are going to train our model on Google Colab.

I am assuming that you already know pretty basics of deep learning computer vision. Before diving into it, make sure you know what’s image classification, what’s transfer learning, how CNN works and some other deep learning fundamentals.

You can check the course. It’s really great.

The steps needed are:

  1. Download image dataset
  2. Load and view your data
  3. Create and train a model
  4. Interpret the results
  5. Make a small web-app out of it

Downloading Image Data

The FastAI library provides functionality for downloading images given a file containing the URLs of these images.

In order to get the URLs, you will have to go to Google Images first and you need to search for your category of choice. Make sure the ad-blocker is turned off or you can try it in incognito mode.

I am going to choose my own categories which are books, watch, earphones, and guitar. You can search for your choices. So what you need to do now is, search your category one by one. Search one and scroll down until enough images are loaded and then execute the following javascript code from the developer console.

Hit Ctrl+Shift+J to open your developer console.

First copy these lines and paste it in your console and press enter. The links of the image file will get downloaded as a file named download. Then you need to rename the file to book.txt or name according to your own preference. Then do repeat the above steps for more classes.

Load and view your data

Open my Google Colab notebook for our further steps. After opening the notebook. Run the first two cells.

This above code in the notebook is for specifying the path and creating its corresponding folder.

After executing the above in the notebook, you need to run the cell below in the notebook and upload the corresponding .txt file.

from google.colab import files
book = files.upload()

You need to upload the .txt file for other classes also.

Now after uploading all the files, You need to execute other cells one by one.

What this code will do is, it will read those URLs which are in the text files and it will download the images in Colab.

After downloading those images, we need to remove the images which can’t be opened.

This code above will do that for us. This verify_images function looks for the images and sees if there’s any problem with it and delete =True will delete it for you. So it’s a really good and easy way to clean our data.

The library has specific data objects called databunches. It helps to read our dataset in a proper folder structure and this helps create validation data randomly. You can read more about in their docs

We have got a data bunch. So you can look inside at the data.classes and you’ll see these are the folders that we created. So it knows that the classes (by classes, we mean all the possible labels) are book, watch, earphones & guitar.

Then you need to run some code below in the notebook. I have commented on the codes so that it would be easier for someone to look into and understand.

Creating a model and initial training

The fastai library lets you create models in a few lines of code and train a CNN model by create_learner method which they provide.

learn = cnn_learner(data, models.resnet34, metrics=error_rate)

This cnn_learner method needs two arguments, the data, and the architecture and it supports many other parameters also.

The created model uses the resnet34 architecture, with weights pretrained on the imagenet dataset.

The fit_one_cycle method is what trains the model. It uses something called the 1 cycle policy, which basically changes the learning rate over time to achieve better results.

Then we save the model. It will take not take more than 5 minutes to train the model.


Lastly, we can use the FastAIs ClassificationInterpretation class to interpret our results.

After running all the cells and training the model. The model file will get exported and you can download the weight file and use it for inference and for production.

You can download the weight file by going to the side by menu and browsing through the folder. After clicking the download, it will get download within some minutes.

Making a small web-app for our image classifier

Clone my Github Repository first by

git clone

After cloning go inside the folder by

cd Custom-Image-Classifier

Then install all the dependencies by

pip3 install -r requirements.txt

After that, inside there will be one folder called models, you need to put your export.pkl file there. This is the path to the folder.


After copying the file you can run the server by

# In Custom-Image-Classifier folder runpython3 app/ serve 

After some seconds it will run on your local server.

Here are some inference results below of my model : –

If you have followed the above steps properly you will be able to create your image classification model easily.

The code covered in this article is available as a Github Repository.

If you have any questions, recommendations or critiques, I can be reached via Twitter or via my mail. Feel free to reach out to me.

References :