import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
import time

# Load YOLO model
model = YOLO("yolov8n.pt")  # Updated to correct model name

# Define trapezoidal restricted area (top-left, top-right, bottom-right, bottom-left)
trapezoid_pts = np.array([[250, 150], [400, 150], [450, 300], [200, 300]], np.int32)

def is_inside_trapezoid(box, trapezoid_pts):
    """Check if the center of a detected object is inside the trapezoidal area."""
    x1, y1, x2, y2 = box
    cx, cy = int((x1 + x2) / 2), int((y1 + y2) / 2)
    return cv2.pointPolygonTest(trapezoid_pts, (cx, cy), False) >= 0

def detect_objects(frame):
    if frame is None:
        return np.zeros((480, 640, 3), dtype=np.uint8), "No input frame"
    
    results = model.predict(frame, conf=0.5)
    annotated_frame = results[0].plot()
    
    # Draw the trapezoid area
    cv2.polylines(annotated_frame, [trapezoid_pts.reshape((-1, 1, 2))], isClosed=True, color=(0, 0, 255), thickness=2)
    
    isAlert = {"alert": [False, ""], "personCount": 0}
    classInIntrusion = ["person", "bicycle", "car", "motorcycle"]
    
    for r in results:
        for box, cls in zip(r.boxes.xyxy, r.boxes.cls):
            class_id = int(cls.item())
            if class_id == 0:  # Person
                isAlert["personCount"] += 1
            if class_id in [0, 1, 2, 3]:  # Person, bicycle, car, motorcycle
                if is_inside_trapezoid(box.tolist(), trapezoid_pts):
                    isAlert["alert"] = [True, classInIntrusion[class_id]]
                    # Mark the intrusion with a red box
                    x1, y1, x2, y2 = map(int, box.tolist())
                    cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 0, 255), 3)
    
    # Add alert text on the frame
    alert_text = f"Intrusion Alert: {isAlert['alert'][0]}, Object: {isAlert['alert'][1]}, Persons: {isAlert['personCount']}"
    cv2.putText(annotated_frame, alert_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    
    return annotated_frame, alert_text

def webcam_feed():
    cap = cv2.VideoCapture(0)
    
    # Check if the webcam opened successfully
    if not cap.isOpened():
        return np.zeros((480, 640, 3), dtype=np.uint8), "Failed to open webcam"
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # Process frame
        result_frame, alert_message = detect_objects(frame)
        
        # Return the processed frame and alert message
        yield result_frame, alert_message

# Create the Gradio interface with webcam
demo = gr.Interface(
    fn=webcam_feed,
    inputs=[],
    outputs=[
        gr.Image(label="Detection Output"),
        gr.Textbox(label="Alert Status")
    ],
    live=True,
    title="YOLO Intrusion Detection",
    description="Real-time detection of persons and vehicles inside a restricted trapezoidal area.",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.queue(max_size=1).launch()