Spaces:
Running
Running
File size: 2,044 Bytes
3e8018c ae0cfcc c5a82b4 3e8018c 0689fbd ae0cfcc 0689fbd 21c63f7 7111c03 ae0cfcc 7111c03 ae0cfcc dd7d891 7764610 dd7d891 0469c68 ae0cfcc 21c63f7 0689fbd ae0cfcc c5a82b4 ae0cfcc 7764610 ae0cfcc c5a82b4 7111c03 |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
import gradio as gr
import cv2
import numpy as np
import tempfile
from datetime import datetime
import shutil
import os # Import os module for file removal
# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
def detect_video(video):
# Check if the video is in bytes (Gradio typically passes it as bytes)
if isinstance(video, bytes):
with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
tmpfile.write(video) # Write the bytes to the temp file
video_path = tmpfile.name # Get the temp file path
else:
video_path = video # In case it's already a file path
cap = cv2.VideoCapture(video_path) # Open the temporary video file
# List to hold results
detection_results = []
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# Get timestamp
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
# Perform detection
results = model(frame)
# Extract bounding boxes and confidence scores
for *xyxy, conf, cls in results.xywh[0]:
x1, y1, x2, y2 = map(int, xyxy) # Convert to integers
detection_results.append({
'timestamp': timestamp,
'coordinates': {'x1': x1, 'y1': y1, 'x2': x2, 'y2': y2},
'confidence': float(conf)
})
cap.release()
# Clean up the temporary file using os.remove
os.remove(video_path)
return detection_results
# Gradio Interface
interface = gr.Interface(fn=detect_video,
inputs=gr.Video(), # Correct input for video
outputs="json",
live=True,
title="YOLOv5 Video Object Detection",
description="Upload a video to detect objects and get bounding boxes with timestamps.")
# Launch the interface with share=True
interface.launch(share=True)
|