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

# Load the model and label encoder (cached for performance)
@st.cache_resource
def load_resources():
    try:
        model = load_model('captains_cv2_model.keras', compile=False)
        with open('label_encoder.pkl', 'rb') as file:
            label_encoder = pickle.load(file)
        return model, label_encoder
    except Exception as e:
        st.error(f"Error loading resources: {str(e)}")
        return None, None

# Preprocess the image (adjusted to match your Colab code)
def preprocess_image(image_path, target_size=(64, 64)):
    try:
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("Failed to load image")
        img = cv2.resize(img, target_size)  # Resize to 64x64
        img_array = np.asarray(img)         # Convert to numpy array, shape: (64, 64, 3)
        img_array = img_array[np.newaxis, :, :, :]  # Add batch dimension, shape: (1, 64, 64, 3)
        return img_array
    except Exception as e:
        st.error(f"Error preprocessing image: {str(e)}")
        return None

# Main Streamlit app
def main():
    # Load model and label encoder
    model, label_encoder = load_resources()
    if model is None or label_encoder is None:
        st.error("Failed to load model or label encoder. Please check the files and try again.")
        return

    # UI setup
    st.title("Image Classification App")
    st.write("Upload an image to get a prediction using the pre-trained CNN model.")

    # File uploader
    uploaded_file = st.file_uploader("Choose an image to classify...", 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)

        # Save the uploaded file temporarily with its original extension
        file_extension = os.path.splitext(uploaded_file.name)[1].lower()
        temp_filename = f"temp_image{file_extension}"
        with open(temp_filename, "wb") as f:
            f.write(uploaded_file.getvalue())

        try:
            # Preprocess the image
            processed_image = preprocess_image(temp_filename)
            if processed_image is None:
                raise ValueError("Image preprocessing failed")

            # Display processed image shape for debugging
            st.write(f"Processed image shape: {processed_image.shape}")

            # Make prediction
            prediction = model.predict(processed_image)
            predicted_class_index = np.argmax(prediction, axis=1)[0]
            predicted_class = label_encoder.inverse_transform([predicted_class_index])[0]

            # Display results
            st.subheader("Prediction Results")
            st.write(f"**Predicted Class:** {predicted_class}")
            st.write(f"**Predicted Class Index:** {predicted_class_index}")
            
            # Display prediction probabilities
            st.write("**Prediction Probabilities:**")
            for class_name, prob in zip(label_encoder.classes_, prediction[0]):
                st.write(f"{class_name}: {prob:.4f}")

        except Exception as e:
            st.error(f"Error during prediction: {str(e)}")

        # Clean up temporary file
        if os.path.exists(temp_filename):
            os.remove(temp_filename)

if __name__ == "__main__":
    main()