import cv2
import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np

# Load the YOLOv5 model (pre-trained on COCO dataset)
model = YOLO('yolov8n.pt')  # You can replace with your custom model if available

def detect_objects(image):
    """
    Detect suspicious objects in the image using YOLO.
    """
    # Convert PIL image to OpenCV format (numpy array)
    image_np = np.array(image)
    results = model(image_np)  # Perform detection

    # Draw bounding boxes on the detected objects
    for result in results:
        boxes = result.boxes  # Bounding boxes
        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])  # Extract box coordinates
            label = box.cls[0]  # Class label
            confidence = box.conf[0]  # Confidence score

            # Draw rectangle and label on the image
            cv2.rectangle(image_np, (x1, y1), (x2, y2), (0, 255, 0), 2)
            text = f"{model.names[int(label)]} ({confidence:.2f})"
            cv2.putText(image_np, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    # Convert back to PIL image
    processed_image = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
    return processed_image

# Create a Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Suspicious Object Detection")
    gr.Markdown("Upload an image or use your webcam to capture one. The app will detect objects using YOLOv5.")

    # Input section
    image_input = gr.Image(type="pil", label="Upload or Capture Image")

    # Output section
    output_image = gr.Image(type="pil", label="Processed Image with Annotations")

    # Button for detection
    detect_button = gr.Button("Detect Suspicious Objects")

    # Link the button to the detection function
    detect_button.click(detect_objects, inputs=[image_input], outputs=[output_image])

demo.launch()