import streamlit as st
from fastai.vision.all import *
from PIL import Image

# Load the model
learn = load_learner('model.pth')
categories = ('granite', 'diorite', 'andacite', 'komatiites', 'gabbro')

# Define the classification function
def classify_image(img):
    pred,idx,probs = learn.predict(img)
    return dict(zip(categories, map(float,probs)))

# Create a file uploader for the image
uploaded_file = st.file_uploader("Choose an rock thin section image...", type=['jpg', 'png'])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Classifying...")
    label = classify_image(image)
    st.write(label)