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

# Load models
model = YOLO("best-3.pt")  # load a custom model for segmentation (protection zone)
model2 = YOLO('yolo11s.pt')  # load a second model for object detection

# Streamlit app title
st.title("Protection Zone and Object Detection")

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

if uploaded_file is not None:
    # Read the image
    image = uploaded_file.read()
    image_np = np.frombuffer(image, np.uint8)
    image_cv = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
    
    # Predict protection zone with the first model
    segment_results = model(image_cv)  # predict segments
    protection_mask = np.zeros(image_cv.shape[:2], dtype=np.uint8)  # create an empty mask
    
    for result in segment_results:
        if result.masks is not None:
            for segment in result.masks.data:
                # Convert segment to numpy array and ensure it's the right shape and type
                segment_array = segment.cpu().numpy().astype(np.uint8)
                segment_array = cv2.resize(segment_array, (image_cv.shape[1], image_cv.shape[0]))
                protection_mask = cv2.bitwise_or(protection_mask, segment_array * 255)
    
    # Create a copy of the original image to draw on
    output_image = image_cv.copy()
    
    # Overlay the protection zone mask on the output image
    protection_overlay = cv2.applyColorMap(protection_mask, cv2.COLORMAP_COOL)
    output_image = cv2.addWeighted(output_image, 0.7, protection_overlay, 0.3, 0)
    
    # Predict objects with the second model
    object_results = model2(image_cv)  # predict objects using model2
    
    for result in object_results:
        boxes = result.boxes.xyxy.cpu().numpy().astype(int)
        for box in boxes:
            x1, y1, x2, y2 = box
            # Check if the object is within the protection zone
            object_mask = np.zeros(image_cv.shape[:2], dtype=np.uint8)
            object_mask[y1:y2, x1:x2] = 1  # create a mask for the object
            # Check overlap
            overlap = cv2.bitwise_and(protection_mask, object_mask)
            color = (0, 0, 255) if np.sum(overlap) > 0 else (0, 255, 0)  # red if in zone, green if outside
            # Draw bounding box around the object
            cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
    
    # Display the final image
    st.image(output_image, caption="Protection Zone and Detected Objects", channels="BGR")
else:
    st.write("Please upload an image to process.")