Sketched Number Prediction App Using Streamlit
Von: Chris Klose
A simple but powerful example of how machine learning models can be implemented in interactive apps build with Streamlit.
If you’re a data scientist or anyone else interested in building machine learning models and embed them into a visual prototype, be sure to check out this tutorial.
In the following we will create a small Streamlit app that predicts numbers drawn by the user on a canvas. This is particularly interesting as it brings together two different disciplines — Machine Learning and Web App Development — in a simple yet elegant fashion.
Let’s get started and see how it works out!
Requirements for the App
Basically two components are required.
- A user interface such as a canvas where numbers can be drawn and the results are to be displayed on.
- A machine learning model that takes an image as input, runs it through the neural network, and gives as back a scalar of what number it is.
Starting off with the User Interface
First we need a web app and a canvas. That’s the easiest part. Streamlit combined with a custom component offers a super simple and fast solution for both. It can be be implemented within minutes and just a few lines of code.
This is the basic skeleton where the canvas component is initialized and the response data is given back to the app as the user draws on it.
Wow, first part is ✅. That was quite simple, right?
Let’s move on to part two where the actual magic is happening.
Image Classification — The Training Data
We are going to use PyTorch and MNIST, a widely known dataset that contains images of numbers from 0–9, drawn with white font color on a black background. It contains 70,000 labelled images with a pixel size of 28x28, hence a total number of 784 features. At that moment you might realize why we chose the same setup for our canvas.
In the next step, a model is constructed and trained.
Getting Advanced — The Model Setup and Training
Since there are already numerous blogposts on how the optimal neural network architecture for image classification models should look like and how to train them, the example model provided by PyTorch will be used.
Our main interest is in applying the trained model. For basic reference, the basic model architecture is shown below. The training script can be found here.
The result is trained model in pickle format that can be integrated into the skeleton created above.
Combining the Web App with the Model
Back in the main app file, we can easily load the saved model into the script by executing the following lines:
PATH = "./mnist_cnn.pt"
model = Net()
model.load_state_dict(torch.load(PATH))The only remaining step is to feed the image into our Neural Network and display the predicted result to the user.
There are a few, preliminary step before doing so. As the image data from Canvas returns RGBA, it has to be converted to greyscale first. We do this with a small helper function:
def rgb2gray(rgb):
return np.dot(rgb[...,:3], [0.299, 0.587, 0.144])Additionally, the size has be reduced from 224x224 as defined in the canvas to 28x28. Also the dimension has to be extended to match the one defined in the model (this part is skipped to keep it simple).
After that, we can finally put the data into the model. The result is a list with the probability for each number. The highest number is considered to be correct. Displaying the number to the user is thanks to Streamlit fairly easy. The adopted lower part of the skeleton looks as follows:
Shout out! That’s it. We are done🎈😎🎉.
Stop! Hold on! There is one more step! Good software engineering includes testing. So lets see how it works and have a little fun in doing so!
You can find a working example of the app and the code:
- https://share.streamlit.io/chrischross/mnistpred/main/app.py
- https://github.com/ChrisChross/MNISTPRed
After a few tries, it turns out our model is very accurate!
As a friend had proven to me, it detects a two even if the format is slightly changed.
Conclusion
In this blog post, we have shown how to develop an interactive web app utilizing a machine learning model for image classification. If you liked it, give it a hands up. If you have any any questions or ideas feel free to contact me as well.
- @ChristianKlose3 on Twitter
- Github ChrisChross
Photo by Nick Hillier on Unsplash
