File size: 1,148 Bytes
6816451 9574f60 6816451 9574f60 6816451 9574f60 6816451 9574f60 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | 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()
|