import streamlit as st
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
from PIL import Image
import os
from huggingface_hub import notebook_login
from huggingface_hub import hf_hub_download

# Title of the Streamlit app
st.title("Yellow Rust Severity Prediction")

authkey= os.getenv('YellowRust')
from huggingface_hub import login 

login(token=authkey)


# Download the model file from Hugging Face
model_path = hf_hub_download(repo_id="shaheer-data/Yellow-Rust-Prediction", filename="final_meta_model.keras")



loaded_model = load_model(model_path)  # Load model using tf.keras directly

# Function to preprocess the uploaded image
def preprocess_image(image):
    # Resize the image to match the model input size (e.g., 224x224 for many pre-trained models)
    image = image.resize((224, 224))  # Adjust size based on your model input
    image = img_to_array(image)       # Convert image to numpy array
    image = image / 255.0            # Normalize pixel values to [0, 1]
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

# Streamlit file uploader
uploaded_file = st.file_uploader("Upload a wheat leaf image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Preprocess the image
    processed_image = preprocess_image(image)

    # Perform prediction
    with st.spinner("Predicting..."):
        prediction = loaded_model.predict(processed_image)
        predicted_class = np.argmax(prediction, axis=1)[0]  # Get the class index
        class_labels = ['0', 'MR', 'MRMS', 'MS', 'R', 'S']  # Update based on your classes
        st.success(f"Predicted Severity Class: {class_labels[predicted_class]}")