import streamlit as st
import tensorflow as tf
from PIL import Image
import os
from huggingface_hub import hf_hub_url, HfApi  # Import Hugging Face utilities

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

# Function to get Hugging Face API token securely
def get_huggingface_api_token():
    if "HUGGINGFACE_API_TOKEN" not in st.session_state:
        st.session_state.HUGGINGFACE_API_TOKEN = st.text_input(
            "Enter your Hugging Face API Token",
            type="password",
            label_visibility="collapsed"
        )
    return st.session_state.HUGGINGFACE_API_TOKEN

# Set Hugging Face token for authentication
huggingface_api_token = get_huggingface_api_token()
if not huggingface_api_token:
    st.warning("Please enter your Hugging Face API Token.")
    st.stop()

HfApi().set_access_token(huggingface_api_token)

# Model repository details
model_repo_id = "shaheer-data/Yellow-Rust-Prediction"
model_file_path = "final_meta_model.keras"

# Construct the model URL
st.write("Loading model from Hugging Face repo:", model_repo_id)
model_url = hf_hub_url(model_repo_id, model_file_path)
loaded_model = tf.keras.models.load_model(model_url)  # Load model using tf.keras directly

# Function to make predictions
def predict_image(image):
    image = image.resize((224, 224))  # Resize to match model input dimensions
    image_array = tf.keras.preprocessing.image.img_to_array(image)
    image_array = tf.expand_dims(image_array, axis=0)  # Expand dimensions for batch prediction
    predictions = loaded_model.predict(image_array)
    return predictions

# Class labels for Yellow Rust severity levels
CLASS_LABELS = [
    "Healthy",
    "Mild Severity",
    "Moderate Severity",
    "Severe Severity",
    "Very Severe",
    "Extreme Severity"
]

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

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

    # Display progress bar
    with st.spinner("Making predictions..."):
        predictions = predict_image(image)
        predicted_class = predictions.argmax(axis=-1)
        st.write(f"Predicted Severity Level: {CLASS_LABELS[predicted_class[0]]} with confidence {predictions[0][predicted_class[0]]:.2f}")

else:
    st.write("Please upload an image file to make predictions.")