import gradio as gr
import cv2
import torch
import numpy as np

# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)

def object_detection(video):
    cap = cv2.VideoCapture(video.name)
    
    while True:
        ret, frame = cap.read()
        
        if not ret:
            break
        
        # Object detection
        results = model(frame)
        
        # Render results
        results.render()
        
        yield results.imgs[0]
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()

iface = gr.Interface(fn=object_detection, inputs=gr.Video(), outputs=gr.Image(label="Object Detection"))
iface.launch()