import streamlit as st
import os
import tensorflow as tf
from tensorflow.keras.models import load_model
import cv2
import numpy as np

# Load the model manually using Keras
model = load_model('Model/keras_model.h5', compile=False)  # Load model without compilation

# Labels directly added to the code
labels = [
    "Nothing",
    "Zip-top cans",
    "Newspaper",
    "Old shoes",
    "Watercolor pen",
    "Disinfectant",
    "Battery",
    "Vegetable leaf",
    "Apple"
]

# Classification dictionary (modified to skip image-based overlays)
classDic = {0: None,
            1: 0,
            2: 0,
            3: 3,
            4: 3,
            5: 1,
            6: 1,
            7: 2,
            8: 2}

# Streamlit page configuration
st.title('Object Detection with Webcam')
st.write("This application uses your webcam to detect objects in real-time.")

# Start the webcam feed for Streamlit
cap = cv2.VideoCapture(0)  # Manually set to camera index 0 (or change to 1, 2 if necessary)
if not cap.isOpened():
    st.error("Error: No camera found!")
    exit()

stframe = st.empty()  # Placeholder for the webcam feed

# Add the 'Exit' button outside the loop so it's only displayed once
exit_button = st.button('Exit', key="exit_button")

# Loop to keep updating webcam feed
while True:
    ret, img = cap.read()
    if not ret or img is None:
        st.error("Error: Failed to capture image!")
        continue

    imgResize = cv2.resize(img, (454, 340))  # Resize captured frame to a specific size

    # Perform prediction using the loaded model
    img_input = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert to RGB
    img_input = cv2.resize(img_input, (224, 224))  # Resize to expected input size for your model
    img_input = img_input / 255.0  # Normalize the image
    prediction = model.predict(tf.convert_to_tensor(img_input[None, ...]))  # Make prediction
    classID = prediction.argmax()  # Get the predicted class ID
    print(f"Predicted class ID: {classID}")

    # Handle the class ID (can be used to trigger different actions based on class)
    if 0 < classID <= len(labels):
        st.write(f"Detected class: {labels[classID - 1]}")
    else:
        st.write(f"Invalid classID: {classID}, setting to default.")

    # Display the resized image in Streamlit
    img_resized = cv2.cvtColor(imgResize, cv2.COLOR_BGR2RGB)  # Convert image to RGB
    stframe.image(img_resized, channels="RGB", caption="Real-time Detection", use_container_width=True)

    # Exit condition based on button press
    if exit_button:
        break

# Release resources and close the window
cap.release()