Create an Image Classification Web App using PyTorch and Streamlit | by Dehao Zhang | Towards Data Science

You just developed a cool ML model.
You are proud of it. You want to show it to your friends through a web demo so they can interact with your model and provide feedback.
However, you are not familiar with common frameworks such as Django and Flask.
You start to ask yourself: Is there a way to build a quick web demo with minimal framework?
What is Streamlit?How to make the UI?How to create the image classification model?What does the result look like?What can we do next?
What is Streamlit?
How to make the UI?
How to create the image classification model?
What does the result look like?
What can we do next?
In this post, I will walk through a quick example of how you can use Streamlit to build a simple web app.
Streamlit is an open-source Python library that makes it easy to build custom web apps for machine learning and data science [1]. Check out its gallery here to see some applications that other people have created.
I have chosen image classification here as an example because computer vision (CV) is one of the most popular areas of AI currently, powered by deep learning algorithms. It also has a wide range of applications, such as classifying medical images to help doctors in disease diagnosis [2]. Learn more about image classification with deep learning models here.
For demonstration purposes, I will use a pretrained ResNet model from PyTorch, and for the same task, you can always use other libraries (TensorFlow, Keras, etc.), other architecture, or even customize your own model.
To see my full Python code, check out my Github page.
Now without further ado, let’s get started!
The first step is to install the Streamlit library, and you can do that using the pip command. I recommend that you use a Python virtual environment to keep your dependencies separately for each project.
$ pip install streamlit
After it is installed successfully, you can do a quick check with a simple ‘Hello World’ app:
$ streamlit hello
For our application, one key element is to enable users to upload images for the model to make predictions, and this can be done with the ‘file_uploader’ function:
import streamlit as stfile_up = st.file_uploader("Upload an image", type="jpg")
Make sure to specify the appropriate file type so it works properly with the classification model.
Then we can display this image:
from PIL import Imageimage = Image.open(file_up)st.image(image, caption='Uploaded Image.', use_column_width=True)
Apart from that, we will also use ‘set_title’ to create a title for the app and ‘write’ to display prediction results.
To use a pretrained model from PyTorch, make sure you have installed both ‘torch’ and ‘torchvision’ library. Now we can create a ResNet model:
from torchvision import models, transformsimport torchresnet = models.resnet101(pretrained=True)
To put it in simple terms, Residual Network (ResNet) is an improved Convolutional Neuron Network (CNN), trained in the well-known ImageNet database, and it has achieved great performance in CV tasks such as image classification and object detection. This particular model ‘resnet101’ means it has 101 layers (Deep!). Read more about ResNet and its variant here.
After we create this model, we need to transform the input image through resizing, normalization, etc. Here we make use of the ‘transforms’ module in ‘torchvision’:
transform = transforms.Compose([    transforms.Resize(256),    transforms.CenterCrop(224),    transforms.ToTensor(),    transforms.Normalize(    mean=[0.485, 0.456, 0.406],    std=[0.229, 0.224, 0.225]    )])
Note that the images are normalized according to its mean and standard deviation. See the documentation for details.
Now we are able to load the image, pre-process it, and make predictions:
img = Image.open(image_path)batch_t = torch.unsqueeze(transform(img), 0)resnet.eval()out = resnet(batch_t)
Next, we can return the top 5 predictions ranked by highest probabilities. Make sure the names of the 1000 categories from ImageNet are available so we can output the predictions through index:
prob = torch.nn.functional.softmax(out, dim=1)[0] * 100_, indices = torch.sort(out, descending=True)return [(classes[idx], prob[idx].item()) for idx in indices[0][:5]]
Now we can run the application by typing ‘streamlit run <the python file containing UI>’ in the command prompt/terminal. Let’s try two examples and see what the results look like.
First example:
Great! Golden retriever is at the top of the list, with high confidence score (~97).
Let’s try another example:
Green snake is at the top of the list! Again, its confidence score is quite high.
One thing to be cautious here is that since the model was trained on ImageNet, its performance might not be as good for images outside of the 1000 categories.
The natural next step after you create a web app is to deploy and host it. As an example, I deployed an Iris Classification app through Heroku and you can check it out through this Github page.
In addition, if you are interested in other neural network models (VGG, Densenet, etc.), you can try upload different images and then compare their performances.
Let’s quickly recap.
We built a simple web app through Streamlit, and in this app, users can upload images to see the model’s top predictions as well as confidence scores. We also went through how to build a pretrained ResNet model through PyTorch.
I hope you enjoy this blog post and please share any thought that you may have :)
Check out my other post on sentiment classification on Yelp reviews using logistic regression: