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
@st.cache_resource
def load_resources():
    # Custom loading to handle compatibility
    try:
        model = load_model('captains_cv2_model.keras', compile=False)  # Load without compiling first
    except Exception as e:
        st.error(f"Model loading failed: {str(e)}")
        raise
    with open('label_encoder.pkl', 'rb') as file:
        le = pickle.load(file)
    return model, le

# Preprocess the image
def preprocess_image(image_path):
    img1 = cv2.imread(image_path)
    img1 = cv2.resize(img1, (64, 64))  # Resize to 64x64
    img1 = np.asarray(img1)  # Shape: (64, 64, 3)
    img1 = img1[np.newaxis, :, :, :]  # Shape: (1, 64, 64, 3)
    return img1

# Main app
def main():
    model, le = load_resources()
    
    st.title("Image Classification App")
    st.write("Upload an image to get a prediction")
    
    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)
        
        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:
            processed_img = preprocess_image(temp_filename)
            st.write(f"Processed image shape: {processed_img.shape}")
            
            prediction = model.predict(processed_img)
            predicted_class = le.inverse_transform([np.argmax(prediction)])
            
            st.write("Prediction:", predicted_class[0])
            st.write("Prediction Probabilities:")
            for class_name, prob in zip(le.classes_, prediction[0]):
                st.write(f"{class_name}: {prob:.4f}")
                
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
        
        if os.path.exists(temp_filename):
            os.remove(temp_filename)

if __name__ == '__main__':
    main()