import os
os.environ["YOLO_CONFIG_DIR"] = "/tmp"  # Prevent permission issues

import torch
import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel  # Required for unpickling

# Patch torch.load() for YOLOv8
torch.serialization.add_safe_globals([DetectionModel])

# Load YOLOv8 model
model = YOLO("yolov8n.pt")

def process_frame(frame):
    img = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    results = model(img)[0]

    for box in results.boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        conf = float(box.conf)
        cls = int(box.cls)
        label = model.names[cls]
        cv2.rectangle(img, (x1, y1), (x2, y2), (0,255,0), 2)
        cv2.putText(img, f"{label} {conf:.2f}", (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)

    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

with gr.Blocks() as demo:
    gr.Markdown("## 🚀 Real-time Object Detection with YOLOv8n")
    webcam = gr.Image(source="webcam", streaming=True)
    output = gr.Image()
    webcam.change(fn=process_frame, inputs=webcam, outputs=output)

demo.launch()