import streamlit as st
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import numpy as np
from PIL import Image
import io

# Load the model
@st.cache_resource
def load_classification_model():
    return load_model('rice_classifier_model_subset.h5')

model = load_classification_model()

# Class names
class_names = ['Arborio', 'Basmati', 'Ipsala', 'Jasmine', 'Karacadag']

st.title('Rice Type Classifier')
st.write('Upload an image of rice, and I\'ll tell you what type it is!')

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file).convert('RGB')  # Convert to RGB
    st.image(image, caption='Uploaded Image', use_column_width=True)

    # Preprocess the image
    img = image.resize((100, 100))
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)
    img_array /= 255.0

    # Make prediction
    predictions = model.predict(img_array)
    predicted_class = class_names[np.argmax(predictions[0])]
    confidence = round(100 * np.max(predictions[0]), 2)

    # Display results
    st.write(f"I think this is **{predicted_class}** rice!")
    st.write(f"Confidence: {confidence}%")

    # Display bar chart of predictions
    st.bar_chart(dict(zip(class_names, predictions[0])))

st.write("---")
st.write("Created with ❤️ by AE")
st.write("Model trained on the [Rice Image Dataset](https://www.kaggle.com/datasets/muratkokludataset/rice-image-dataset/)")