The world’s leading publication for data science, AI, and ML professionals.

Deploying an Image Classification Web App with Python

How to deploy your machine learning model using Streamlit and Heroku

It is no doubt that doing a Data Science and machine learning project, starting from collecting the data, processing the data, visualizing insights about the data, and developing a machine learning model to do a predictive task is a fun thing to do. What makes it more fun and doable is that we can do all of those steps in our local machine and then be done with it.

However, wouldn’t it be awesome if other people can make use of our machine learning model to do fun and cool stuff? The true magic of machine learning comes when our model can get into other people’s hand and they can do useful stuff from it.

Creating a web app is one of the solutions such that other people can make use of our machine learning model. Fortunately, it is very simple to create a web app nowadays.

If you’re using Python, you can use Streamlit library to create a simple machine learning web app in your local machine. To deploy the web app to be accessible to other people, then we can use Heroku or other cloud platforms. In this article, I will show you step-by-step on how to create your own simple web app for Image Classification using Python, Streamlit, and Heroku.

If you haven’t installed Streamlit yet, you can install it by running the following pip command in your prompt.

pip install streamlit

Image Classification Data

For this image classification example, the rock-paper-scissor dataset created by Laurence Moroney will be used. You can download the data on his website. As a summary, rock-paper-scissor is a synthetic dataset which contains people’s hand forming either rock, paper, or scissor shape. In total, there are 2520 training images and 372 test images.

Below is the sneak-peek of what the synthetic image data look like:

Sneak-peek of training images of rock-paper-scissor dataset
Sneak-peek of training images of rock-paper-scissor dataset

Overall there are three steps will be covered in this article:

  • Creating the first Python file to load the data, build the model, and finally train the model.
  • Creating the second Python file to load the model and to build the web app.
  • Deploying the web app using Heroku.

Load the Data, Build the Model, and Train the Model

As a first step, download the training set and the test set and save it to the directory of your choice. Next, you need to unzip them. There will be two folders, one called ‘rps’ for training images and the other called ‘rps-test-set’ for test images. After that, create a Python file called rps_model.py to load the data, build the Machine Learning model, and train the model.

Load the Data

Before building the model, you need to first specify the path to the training set folder and the test set folder in your local machine using os library.

import os
train_dir = os.path.join('directory/to/your/rps/')
test_dir = os.path.join('directory/to/your/rps-test-set/')

Next, we can use image generator from TensorFlow library to generate training and test set as well as to automatically label the data. Another advantage of using image generator is that we can do data augmentation ‘on the fly’ to increase the number of training set by zooming, rotating, or shifting the training images. Also, we can split a certain portion of the training set for the validation set to compare the performance of the model.

from tensorflow.keras.preprocessing.image import ImageDataGenerator

def image_gen_w_aug(train_parent_directory, test_parent_directory):

    train_datagen = ImageDataGenerator(rescale=1/255,
                                      rotation_range = 30,  
                                      zoom_range = 0.2, 
                                      width_shift_range=0.1,  
                                      height_shift_range=0.1,
                                      validation_split = 0.15)

    test_datagen = ImageDataGenerator(rescale=1/255)

    train_generator =          train_datagen.flow_from_directory(train_parent_directory,
                                  target_size = (75,75),
                                  batch_size = 214,
                                  class_mode = 'categorical',
                                  subset='training')

    val_generator = train_datagen.flow_from_directory(train_parent_directory,
                                  target_size = (75,75),
                                  batch_size = 37,
                                  class_mode = 'categorical',
                                  subset = 'validation')

    test_generator = test_datagen.flow_from_directory(test_parent_directory,
                                 target_size=(75,75),
                                 batch_size = 37,
                                 class_mode = 'categorical')
    return train_generator, val_generator, test_generator

After creating the function above, now you can call the function and pass the argument with train_dir and test_dir variables that you have defined before.

train_generator, validation_generator, test_generator = image_gen_w_aug(train_dir, test_dir)

And with this step, you have already loaded the data!


Build the Machine Learning Model

As a machine learning model, you can build your own CNN model if you wish. But in this article, transfer learning method will be applied instead. The InceptionV3 model with pre-trained weights from ImageNet is used.

However, since this model is very deep, the model used will be until the layer called mixed_5 and all the weights up until this layer are fixed to speed up the training time. From this layer, the model will be flattened to be connected to a dense layer and then to the final dense output. All the weights after mixed_5 are trainable.

from tensorflow.keras.applications.inception_v3 import InceptionV3
from tensorflow.keras.layers import Flatten, Dense, Dropout
def model_output_for_TL (pre_trained_model, last_output):    
    x = Flatten()(last_output)

    # Dense hidden layer
    x = Dense(512, activation='relu')(x)
    x = Dropout(0.2)(x)

    # Output neuron. 
    x = Dense(3, activation='softmax')(x)

    model = Model(pre_trained_model.input, x)

    return model
pre_trained_model = InceptionV3(input_shape = (75, 75, 3), 
                                include_top = False, 
                                weights = 'imagenet')
for layer in pre_trained_model.layers:
  layer.trainable = False
last_layer = pre_trained_model.get_layer('mixed5')
last_output = last_layer.output
model_TL = model_output_for_TL(pre_trained_model, last_output)

If you haven’t imported InceptionV3 model yet, it will take a couple of minutes to download the model.

At this step, you’ve built the model! Next, we need to train the model.


Train and Save the Machine Learning Model

Now it’s time for us to train the model. Before you train the model, first you need to compile the model. In this article, the optimizer will be Adam and since we have a classification problem, then we should use categorical cross-entropy as the loss. For the metrics, we will use accuracy.

After compiling the model, now we can train the model. After the model is trained, then we need to save the trained model. You should then have a new hdf5 file called my_model.hdf5 in the same directory as your Python file.

model_TL.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
history_TL = model_TL.fit(
      train_generator,
      steps_per_epoch=10,  
      epochs=20,
      verbose=1,
      validation_data = validation_generator)
tf.keras.models.save_model(model_TL,'my_model.hdf5')

Now you’re done with the modelling process. You can save and close the Python file.


Create a Streamlit Web App

To create a web app with Streamlit, first thing that we need to do is creating a new Python file, let’s call it rps_app.py . In this Python file, first we need to load the trained model that we have saved before.

import tensorflow as tf
model = tf.keras.models.load_model('my_model.hdf5')

The next step is to write a header and any other texts that you want to put into your web app. To write a header, simply use write attribute from Streamlit and then put # before your text. To write a simple text, we can also use write method and then proceed with your text without adding #

To let the user upload their own image to your web app, simply use file_uploader attribute from Streamlit library. For more detailed information regarding other options available on Streamlit, you can check it out at Streamlit API documentation page.

import streamlit as st
st.write("""
         # Rock-Paper-Scissor Hand Sign Prediction
         """
         )
st.write("This is a simple image classification web app to predict rock-paper-scissor hand sign")
file = st.file_uploader("Please upload an image file", type=["jpg", "png"])

The next important step is to process the image the user has uploaded. The processing step including resizing the image to the same size as training and validation images. After resizing the image, then the loaded model should predict in which category this image belongs.

import cv2
from PIL import Image, ImageOps
import numpy as np
def import_and_predict(image_data, model):

        size = (150,150)    
        image = ImageOps.fit(image_data, size, Image.ANTIALIAS)
        image = np.asarray(image)
        img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        img_resize = (cv2.resize(img, dsize=(75, 75),    interpolation=cv2.INTER_CUBIC))/255.

        img_reshape = img_resize[np.newaxis,...]

        prediction = model.predict(img_reshape)

        return prediction
if file is None:
    st.text("Please upload an image file")
else:
    image = Image.open(file)
    st.image(image, use_column_width=True)
    prediction = import_and_predict(image, model)

    if np.argmax(prediction) == 0:
        st.write("It is a paper!")
    elif np.argmax(prediction) == 1:
        st.write("It is a rock!")
    else:
        st.write("It is a scissor!")

    st.text("Probability (0: Paper, 1: Rock, 2: Scissor")
    st.write(prediction)

After that, you need to save the Python file in the same directory as your previous Python file.

We basically all set right now! To check what our web app looks like, open your prompt, and then navigate to the working directory of your Python files. In the working directory, you can type the following command:

streamlit run rps_app.py

Now you will see from your prompt that you can check your web app on your localhost. If you wait a little bit, a new window will be launched shortly after you run your Streamlit app. Below is the screenshot of the simple image classification web app.


Deploy Your Web App with Heroku

So far, you’ve built the web app locally on your computer. In order for other people to be able to use your web app, you can utilize Heroku.

Additional Files for Deployment

Before we deploy the web app, we need to create three additional files in addition to Python files that we have created to build the app. These three files are:

  • requirements.txt: this is the text file that we need to create to tell Heroku to install the necessary Python packages needed to deploy our machine learning model. In this tutorial, we used four Python libraries to build the app: numpy, streamlit, tensorflow, and pillow. Hence, we need to specify the relevant version of those libraries in this text file.
tensorflow==2.0.0
streamlit==0.62.0
numpy==1.16.5
pillow==6.2.0

If you are unsure about the version of the Python libraries that you are using, you can use ‘version attribute in your Python environment or type ‘conda list’ in your conda prompt.

  • setup.sh: this file is necessary to handle the server and port number of our app on Heroku.
mkdir -p ~/.streamlit/                                               echo "                       
[server]n                       
port = $PORTn                       
enableCORS = falsen                       
headless = truen                       
n                       
" > ~/.streamlit/config.toml
  • Procfile: this is the file of your configuration to tell Heroku how and which files to be executed.
web: sh setup.sh && streamlit run rps_app.py

Now that you have these three files, place them in the same directory as your Python files.

Next, let’s jump into Heroku.

Create Your Heroku Account

If you already have Heroku account, you can skip this step. If you haven’t, then you can directly go to Heroku. There you’ll find a ‘Sign Up’ button, click that button and fill the necessary information. After that, you need to confirm your new Heroku account with your E-Mail.

Deploy Your ML Model with Heroku Git

After you have created a Heroku account, the next thing that you need to do is to install Heroku CLI (command-line interface). Next, open your command prompt and type the following command.

heroku login

This command will direct you to login with your Heroku account.

Next, move to the directory of your app files with the command prompt. In the directory of your app files, type the following command.

heroku create your_app_name

You can change the your_app_name with your preference as the app name. Next, you need to initiate an empty git repository by typing the following command.

git init

Next, you want to submit all of your app files to the empty repository that you’ve just created. To do this, you can type git add . command and then you need to push them by typing commit and push commands as follows.

git add .
git commit -m "your message"
git push heroku master

Now your app files will be deployed to Heroku with Heroku Git and after the process is completed, you will see the URL of your web app that can be accessed via the internet. The format of your web app URL is something like _your_appname.herokuapp.com.

Below is the example of a simple web app that has been deployed with Heroku.

And that’s it! Your image classification web app has been deployed!

If you want to see the code of the web app or other necessary files in this tutorial, you can see it on my GitHub.


Related Articles