hikmatfarhat/MNIST

1import streamlit as st2from streamlit_image_select import image_select3from PIL import Image4import pandas as pd5from transformers import AutoModel6import torch7from torchvision.datasets import MNIST8from torchvision.transforms import ToTensor9import numpy as np10st.title("MNIST classifier")11st.markdown("This is a simple MNIST classifier using a simple neural network")12st.markdown("select a digit from the sidebar and the classifier will give you the probability of each digit")13classifier=AutoModel.from_pretrained("hikmatfarhat/MNIST_Classifier",trust_remote_code=True)14softmax=torch.nn.Softmax(dim=1)15np.set_printoptions(precision=3,suppress=True)16imgs=[]17 18for i in range(15):19    imgs.append(Image.open(f"img{i}.png"))20with st.sidebar:21    image=image_select("select a digit",imgs,use_container_width=False)22image=ToTensor()(image)23output=classifier(image)24output=torch.squeeze(softmax(output))25output=output.detach().numpy()26## streamlit doesn't show the proper label for the column name when using numpy array27## so we convert it to pandas dataframe28df=pd.DataFrame(output,columns=["prob"])29st.dataframe(df,column_config={"prob":st.column_config.Column(30    "Probability",required=True)}31    ,hide_index=False32    )33