import gradio as gr
import cv2
import torch

def detect_objects(frame):
    results = model(frame)

    for *box, conf, cls in results.xyxy[0]:
        label = f'{model.names[int(cls)]} {conf:.2f}'
        frame = cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 0, 0), 2)
        frame = cv2.putText(frame, label, (int(box[0]), int(box[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
video = gr.inputs.Video(type='opencv', label="输入视频")
output = gr.outputs.Image(type='pil',label="检测结果")

iface = gr.Interface(fn=detect_objects, inputs=video, outputs=output, title="目标检测")
iface.launch(share=True)