import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
from io import BytesIO
import requests
import plotly.express as px

# Set page configuration
st.set_page_config(page_title="Scene Classifier", layout="wide")

# Load the model
@st.cache_resource
def load_model():
    try:
        return tf.keras.models.load_model('epoch_26.h5')
    except:
        st.error("Model file not found. Please make sure 'model_epoch_11.h5' is in the same directory as this script.")
        return None

model = load_model()

# Define class names based on your dataset
class_names = ['buildings', 'forest', 'glacier', 'mountain','human', 'sea', 'street']

def preprocess_image(img):
    """Preprocess image for prediction"""
    img = img.convert('RGB')
    img = img.resize((224, 224))
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)
    img_array = img_array / 255.0  # Normalize the image
    return img_array


def predict_and_display(img):
    """Make prediction and display results with confidence scores"""
    # Preprocess the image
    img_array = preprocess_image(img)

    # Make prediction
    predictions = model.predict(img_array)[0]  # Get the first (and only) prediction

    # Get predicted class and confidence
    predicted_index = np.argmax(predictions)
    predicted_class = class_names[predicted_index]
    confidence_score = predictions[predicted_index] * 100  # Convert to percentage

    # Display predicted label and confidence
    st.markdown(f"### **Predicted:** {predicted_class} ({confidence_score:.2f}%)")
    st.markdown(f"🧠 Model is **{confidence_score:.2f}%** confident that the image is a **{predicted_class}**.")
    # Create a dataframe for plotting
    confidence_data = {"Category": class_names, "Confidence (%)": [score * 100 for score in predictions]}
    
    # Create a bar chart
    fig = px.bar(
        confidence_data, 
        x="Confidence (%)", 
        y="Category", 
        orientation="h",
        text_auto=".2f",
        title="Confidence Scores",
        color="Category",
        color_discrete_sequence=px.colors.qualitative.Set1
    )
    fig.update_traces(textposition="outside")  # Show confidence values outside bars
    fig.update_layout(
        yaxis={"categoryorder": "total ascending"},  # Sort by confidence
        height=600,  # Increase chart height
        width=700,  # Increase chart width
        margin=dict(l=30, r=30, t=50, b=50)  # Adjust margins for better spacing
    )

    # Display image & graph side by side
    col1, col2 = st.columns([1, 1])
    with col1:
        img = img.resize((512, 512))
        st.image(img)
    with col2:
        st.plotly_chart(fig, use_container_width=True)




# Streamlit app UI
st.title("🌄 Scene Classifier")
st.markdown("""
This app classifies images into one of the following categories:
- 🏢 Buildings    🌲 Forest    ❄️ Glacier    ⛰️ Mountain    🌊 Sea    🛣️ Street    🚹 Human
""")

# Create tabs for URL input and file upload
tab1, tab2 = st.tabs(["Upload Image", "Enter URL"])

with tab1:
    uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        try:
            img = Image.open(uploaded_file)
            if st.button("Classify Uploaded Image"):
                predict_and_display(img)
        except Exception as e:
            st.error(f"Error processing the image: {e}")

with tab2:
    image_url = st.text_input("Enter Image URL:")
    if st.button("Classify from URL") and image_url:
        try:
            response = requests.get(image_url)
            if response.status_code == 200:
                img = Image.open(BytesIO(response.content))
                predict_and_display(img)
            else:
                st.error(f"Error fetching the image. Status code: {response.status_code}")
        except Exception as e:
            st.error(f"Error processing the URL: {e}")