import streamlit as st
import keras
import cv2
import numpy as np
import pickle
from PIL import Image

st.set_page_config(page_title="CV2 Image Detection", layout="centered")

st.title("🖼️ A Minor Project on Image Detection using CNN and CV2 (A Simplified Clone of CNN Explainer)")

# Cache the model and label encoder
@st.cache_resource
def load_model():
    try:
        model_path = r"cv2_model.keras"
        label_encoder_path = r"label_encoder.pkl"

        model = keras.models.load_model(model_path)
        with open(label_encoder_path, 'rb') as file:
            label_encoder = pickle.load(file)

        return model, label_encoder
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None, None

# File uploader for image input
user_input_img = st.file_uploader("📤 Upload an Image", type=["jpg", "jpeg", "png"])

# Initialize session state for prediction
if "prediction" not in st.session_state:
    st.session_state["prediction"] = None
    st.session_state["uploaded_img"] = None

if user_input_img is not None:
    # Load and preprocess image
    img = Image.open(user_input_img)
    img_array = np.array(img)
    img_resized = cv2.resize(img_array, (64, 64))
    img_resized = np.expand_dims(img_resized, axis=0)  # Another way to batch dimension instead of np.newaxis

    # Store image in session state
    st.session_state["uploaded_img"] = img

# Button for Prediction
if st.button("🔍 Predict"):
    if st.session_state['uploaded_img'] is None:
        st.warning("⚠️ Please upload an image before predicting!")
        
    else:
        model, encoder = load_model()
        if model and encoder:
            # Make prediction
            predicted_label = encoder.inverse_transform(np.argmax(model.predict(img_resized), axis = 1))

            # Store prediction in session state
            st.session_state["prediction"] = predicted_label

# Display result in two columns if prediction exists
if st.session_state["prediction"]:
    col1, col2 = st.columns([1, 1])

    with col1:
        st.success(f"**Prediction:** {st.session_state['prediction']}")

    with col2:
        st.image(st.session_state["uploaded_img"], caption="Uploaded Image", use_container_width =True)