Skip to content

Contact sales

By filling out this form and clicking submit, you acknowledge our privacy policy.

Deploying Image Classification on the Web with Streamlit and Heroku

Jun 8, 2020 • 8 Minute Read

Introduction

In this guide, you will learn to deploy an image classifier on the web using Streamlit and Heroku. An image classifier is a computer vision algorithm that is able to assign an image to a particular predefined class based on the content of the image file. It is used to enable applications to perform classification and sorting tasks on image data.

For this guide, you will download a simple binary classifier from Google Teachable Machine. This guide assumes you have intermediate knowledge of Python. The other prerequisites for this guide are:

  1. Basic knowledge of Streamlit

  2. Basic knowledge of image classification in machine learning

  3. Basic knowledge of Google's Teachable Machine platform

Here's a real-world scenario:

Imagine you are a machine learning engineer at a remote-first AI/ML startup. You are part of the research and development team that specializes in AI and ML. Being remote first, you need to share your image classification app idea with not only your team members, but also other teams and company staff that may not be savvy about ML but want to test the practicality of your solution. You will also need to roll out your image classification solution to the general public. This means that it should be simple, intuitive, and easily accessible, and that you should know how to deploy ML solutions.

Environment Setup

To get your environment ready, install the relevant packages by running the command

pip install keras streamlit pillow numpy

With the four packages, you are ready to start a simple image classifier.

Image Classification with Google Teachable Machine

Teachable Machine allows you to train classifier models on the web. You can create an image, audio, or pose project. In this case, choose an image classification project. For this guide, you will use the brain tumor detection dataset on Kaggle. Follow the link to download the dataset.

Once on the image classification page, label Class 1 as Brain Tumor and Class 2 as No Brain Tumor. Upload the images to the appropriate class. After uploading the data, click the Train Model button and wait for results.

Once done, you will be able to download the weights file with a .h5 extension. This is what you will use for your classification task.

Streamlit-Based Image Classifier

In your development folder, create a file named app.py. This is the file that holds the Streamlit code responsible for displaying content on the webpage.

To set up the page by giving it a header, title, and description, write the following code.

      import streamlit as st
st.title("Image Classification with Google's Teachable Machine")
st.header("Brain Tumor MRI Classification Example")
st.text("Upload a brain MRI Image for image classification as tumor or no-tumor")

Next, handle the file upload, processing by the classifier from Teachable Machine, and finally, displaying results. To ensure neat code and follow best practice, create a separate file to handle the actual classification.

Image Classification Code

In the same folder that holds app.py, create a file called img_classification.py. This will hold the classification function that you will call in the app.py file for image classification.

To perform classification, load the weights file that you downloaded in .h5 format into a keras model.

      import keras
from PIL import Image, ImageOps
import numpy as np


def teachable_machine_classification(img, weights_file):
    # Load the model
    model = keras.models.load_model(weights_file)

    # Create the array of the right shape to feed into the keras model
    data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)
    image = img
    #image sizing
    size = (224, 224)
    image = ImageOps.fit(image, size, Image.ANTIALIAS)

    #turn the image into a numpy array
    image_array = np.asarray(image)
    # Normalize the image
    normalized_image_array = (image_array.astype(np.float32) / 127.0) - 1

    # Load the image into the array
    data[0] = normalized_image_array

    # run the inference
    prediction = model.predict(data)
    return np.argmax(prediction) # return position of the highest probability

Streamlit File Handling and Results

This part of the code performs image upload handling, running inference from the function in the img_classification.py file and displaying the results. Just below the imports in app.py, add this line to allow you to use your classification function.

      from img_classification import teachable_machine_classification

The below code will be added to your app.py file.

      uploaded_file = st.file_uploader("Choose a brain MRI ...", type="jpg")
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded MRI.', use_column_width=True)
        st.write("")
        st.write("Classifying...")
        label = teachable_machine_classification(image, 'brain_tumor_classification.h5')
        if label == 0:
            st.write("The MRI scan has a brain tumor")
        else:
            st.write("The MRI scan is healthy")

To run your application on localhost, open up your terminal or shell, navigate to the current working directory, and run the command

streamlit run app.py

The app will run locally and be available via the URL https://localhost:8501

Heroku Account Setup

At this point, your app is set up and running locally. It's time to share it. Register for a Heroku account to allow you to deploy your awesome classifier to the web. After registration, download the Heroku CLI.

Deploying the Solution

Generate a requirements.txt file that will hold the libraries used and their versions. You can use the pipreqs command. The file should look something like this:

      numpy==1.16.4
streamlit==0.52.1
pillow
keras
tensorflow==2.0.0b1

Create a setup.sh file and a procfile file. These files will instruct Heroku on what to do to set up the app and get it running. In the setup.sh file, write the below code, which will will create a streamlit folder with a credentials.toml and a config.toml file.

      mkdir -p ~/.streamlit/
echo "\
[general]\n\
email = \"[email protected]\"\n\
" > ~/.streamlit/credentials.toml
echo "\
[server]\n\
headless = true\n\
enableCORS=false\n\
port = $PORT\n\
" > ~/.streamlit/config.toml

Next, the procfile will be the one that executes setup.sh. Then run the Streamlit app.py file procfile code:

      web: sh setup.sh && streamlit run app.py

Initiate an empty Git repository using the command git init.

In your terminal, navigate to the code's working directory and log in to Heroku using the CLI command heroku login. To deploy, run the command heroku create. Upon completion of this command, Heroku will assign an app name and URL to your app, which will allow you to access it via the web.

Last, push your code to your Heroku instance using the Git commands below.

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

To check if deployment was successful, run the command heroku ps scale:web=1.

Conclusion

You have now learned to deploy machine learning solutions on the web using Streamlit to build the interface and Heroku to serve the app on the web. These skills are vital for real world roles such as ML engineer, ML devops, software engineer with an interest in ML, and freelance ML enthusiasts who would like to share their work. For a more in-depth look at Heroku, consider this tutorial.

You can also further build on the skills in this guide by researching cloud-based ML deployment solutions such as Amazon AWS, Microsoft Azure, GCP, and FloydHub.