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

# Set page configuration
st.set_page_config(
    page_title="Image Detection App",
    page_icon="📸",
    layout="centered",
    initial_sidebar_state="expanded"
)

# Load the trained model and label encoder
@st.cache_resource
def load_resources():
    model = load_model("captains_cv2_model.keras")
    with open("label_encoder.pkl", "rb") as f:
        le = pickle.load(f)
    return model, le

model, label_encoder = load_resources()

# Function to preprocess the uploaded image
def preprocess_image(uploaded_file):
    # Save the uploaded file temporarily to disk
    temp_path = "temp_image.jpg"
    with open(temp_path, "wb") as f:
        f.write(uploaded_file.read())
    
    # Read the image using cv2.imread
    img = cv2.imread(temp_path)
    # Resize to the model's expected input size (64, 64)
    img = cv2.resize(img, (64, 64))  # Note: cv2.resize takes (width, height), not (height, width, channels)
    # Add new axis for batch dimension
    img = img[np.newaxis, :, :, :]
    
    # Clean up the temporary file
    os.remove(temp_path)
    return img

# Sidebar
st.sidebar.title("About")
st.sidebar.info(
    "This app uses a Convolutional Neural Network (CNN) to classify images into one of 10 categories. "
    "Upload an image, and the model will predict its class!"
)
st.sidebar.markdown("### Classes")
st.sidebar.write(
    "The model can predict: lifeboat, ladybug, pizza, bell pepper, school bus, koala, espresso, red panda, orange, sports car."
)

# Main content
st.title("📸 Image Classification App")
st.markdown("Upload an image below, and let the model predict its class!")

# File uploader
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)
    uploaded_file.seek(0)  # Reset file pointer after reading for display
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Preprocess the image
    processed_image = preprocess_image(uploaded_file)

    # Make prediction
    with st.spinner("Predicting..."):
        # Predict and decode as per your specified steps
        prediction = model.predict(processed_image)
        predicted_class_idx = np.argmax(prediction, axis=1)[0]
        predicted_class = label_encoder.inverse_transform([predicted_class_idx])[0]

    # Display the prediction
    st.success("Prediction complete!")
    st.markdown(f"### Predicted Class: **{predicted_class}**")
    st.write(f"Prediction Confidence: {prediction[0][predicted_class_idx]:.4f}")

    # Optional: Display confidence scores for all classes
    if st.checkbox("Show confidence scores for all classes"):
        class_names = label_encoder.classes_
        confidence_scores = {class_names[i]: float(prediction[0][i]) for i in range(len(class_names))}
        st.bar_chart(confidence_scores)

else:
    st.info("Please upload an image to get started.")

# Footer
st.markdown("---")
st.markdown("Created with ❤️ using Streamlit | Hosted on [Hugging Face Spaces](https://huggingface.co/spaces)")