import streamlit as st
from streamlit_image_select import image_select
from PIL import Image
import pandas as pd
from transformers import AutoModel
import torch
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import numpy as np
st.title("MNIST classifier")
st.markdown("This is a simple MNIST classifier using a simple neural network")
st.markdown("select a digit from the sidebar and the classifier will give you the probability of each digit")
classifier=AutoModel.from_pretrained("hikmatfarhat/MNIST_Classifier",trust_remote_code=True)
softmax=torch.nn.Softmax(dim=1)
np.set_printoptions(precision=3,suppress=True)
imgs=[]

for i in range(15):
    imgs.append(Image.open(f"img{i}.png"))
with st.sidebar:
    image=image_select("select a digit",imgs,use_container_width=False)
image=ToTensor()(image)
output=classifier(image)
output=torch.squeeze(softmax(output))
output=output.detach().numpy()
## streamlit doesn't show the proper label for the column name when using numpy array
## so we convert it to pandas dataframe
df=pd.DataFrame(output,columns=["prob"])
st.dataframe(df,column_config={"prob":st.column_config.Column(
    "Probability",required=True)}
    ,hide_index=False
    )