import streamlit as st
from PIL import Image
import numpy as np
import cv2
from ultralytics import YOLO

# Load your PPE detection model
model = YOLO('ppe-kit.pt')  # Replace with your PPE detection model path

# Set page configuration
st.set_page_config(
    page_title="PPE Kit Detection",
    page_icon="🦺",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Sidebar contents
st.sidebar.title("🦺 PPE Kit Detection")
st.sidebar.write("Upload an image to detect PPE kits.")

# Upload image
uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

# Main content
st.title("🦺 PPE Kit Detection")
st.write(
    """
    This application uses a state-of-the-art YOLO model to detect PPE kits and safety measures in images.
    It identifies items like helmets, masks, safety vests, and more. Upload an image to see the detection in action.
    """
)

if uploaded_file is not None:
    # Display uploaded image
    image = Image.open(uploaded_file)
    
    # Create two columns for displaying the image
    col1, col2 = st.columns(2)

    with col1:
        st.header("Uploaded Image")
        st.image(image, caption='Uploaded Image', use_column_width=True)
    
    with col2:
        st.header("Detected PPE and Safety Items")

        # Convert the image to OpenCV format
        image_cv = np.array(image)
        image_cv = image_cv[:, :, ::-1].copy()  # Convert RGB to BGR

        # Run inference
        results = model(image_cv)

        # Process the results
        boxes = results[0].boxes
        labels = {
            0: 'Helmet', 1: 'Mask', 2: 'NO-Helmet', 3: 'NO-Mask', 4: 'NO-Safety Vest',
            5: 'Person', 6: 'Safety Cone', 7: 'Safety Vest', 8: 'machinery', 9: 'vehicle'
        }

        detected_items = []  # List to store detected items and accuracy

        # Define a set of class IDs to skip (machinery, vehicle, and person)
        skip_classes = {5, 8, 9}  # Person, machinery, vehicle

        # Draw bounding boxes and labels on the image
        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            confidence = box.conf[0]
            class_id = int(box.cls[0])
            label = labels.get(class_id, 'Unknown')

            # Skip drawing and storing items if they belong to the skip_classes set
            if class_id in skip_classes:
                continue

            # Store detected items with accuracy
            detected_items.append(f"{label} ({confidence * 100:.2f}%)")

            # Draw the box
            cv2.rectangle(image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
            # Put the label
            cv2.putText(image_cv, f"{label} {confidence:.2f}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)

        # Convert the image back to RGB
        image_cv = image_cv[:, :, ::-1].copy()
        st.image(image_cv, caption='Detected PPE and Safety Items', use_column_width=True)

    # Centered detected items below the image
    st.markdown("<h2 style='text-align: center;'>Detected Items with Accuracy:</h2>", unsafe_allow_html=True)
    
    # Display the detected items in a centered fashion
    for item in detected_items:
        st.markdown(f"<h4 style='text-align: center;'>{item}</h4>", unsafe_allow_html=True)

    st.success("Detection complete!")

else:
    st.info("Please upload an image to start detection.")

# Footer
st.sidebar.markdown("### About")
st.sidebar.info(
    "This app uses a YOLO model to detect PPE kits like helmets, masks, safety vests, and other safety-related items. "
    "Upload an image and the model will identify these items in the image."
)