import streamlit as st
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from PIL import Image
import os

def load_model_from_huggingface():
    """Loads the model from Hugging Face Hub."""
    from huggingface_hub import from_pretrained_keras
    model = from_pretrained_keras("shaheer-data/Yellow-Rust-Prediction")
    return model

def preprocess_image(image):
    """Preprocesses the uploaded image for model prediction."""
    image = image.resize((224, 224))  # Assuming input size of 224x224 for the model
    image_array = np.array(image)
    image_array = image_array / 255.0  # Normalize pixel values to [0, 1]
    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
    return image_array

def predict_severity(model, image):
    """Predicts the severity using the model."""
    predictions = model.predict(image)
    class_names = ['0', 'MR', 'MRMS', 'MS', 'R', 'S']
    predicted_class = class_names[np.argmax(predictions)]
    return predicted_class, predictions

# Streamlit App
st.title("Disease Severity Prediction App")
st.write("Upload an image to predict the severity of the disease.")

# Image Upload
uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_image:
    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)

    # Load model
    with st.spinner("Loading model..."):
        model = load_model_from_huggingface()

    # Preprocess image
    image = Image.open(uploaded_image)
    preprocessed_image = preprocess_image(image)

    # Predict severity
    with st.spinner("Predicting severity..."):
        predicted_class, prediction_scores = predict_severity(model, preprocessed_image)

    # Display results
    st.success(f"Predicted Class: {predicted_class}")
    st.write("Prediction Scores:")
    for class_name, score in zip(['0', 'MR', 'MRMS', 'MS', 'R', 'S'], prediction_scores[0]):
        st.write(f"{class_name}: {score:.4f}")