Spaces:
Sleeping
Sleeping
File size: 3,578 Bytes
caff61e 359afbb ae0ee90 359afbb 054b852 fa9a701 ae0ee90 359afbb ae0ee90 5b68bf2 6de980c ae0ee90 054b852 ae0ee90 5b68bf2 ae0ee90 5b68bf2 ae0ee90 5b68bf2 ae0ee90 5b68bf2 ae0ee90 5b68bf2 ae0ee90 5b68bf2 ae0ee90 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import torch
import numpy as np
import gradio as gr
import cv2
import time
import os
from pathlib import Path
# Create cache directory for models
os.makedirs("models", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load YOLOv5 Model
model_path = Path("models/yolov5n.pt")
if model_path.exists():
print(f"Loading model from cache: {model_path}")
model = torch.hub.load("ultralytics/yolov5", "custom", path=str(model_path), source="local").to(device)
else:
print("Downloading YOLOv5n model and caching...")
model = torch.hub.load("ultralytics/yolov5", "yolov5n", pretrained=True).to(device)
torch.save(model.state_dict(), model_path)
# Configure model
model.conf = 0.5
model.iou = 0.5
model.classes = None
if device.type == "cuda":
model.half()
else:
torch.set_num_threads(os.cpu_count())
model.eval()
# Generate colors for bounding boxes
np.random.seed(42)
colors = np.random.uniform(0, 255, size=(len(model.names), 3))
def detect_objects(image):
if image is None:
return None
output_image = image.copy()
results = model(image, size=640)
detections = results.pred[0].cpu().numpy()
for *xyxy, conf, cls in detections:
x1, y1, x2, y2 = map(int, xyxy)
class_id = int(cls)
color = colors[class_id].tolist()
cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
label = f"{model.names[class_id]} {conf:.2f}"
cv2.putText(output_image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
return output_image
def process_video(video_path):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return "Error: Could not open video file."
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
fps = cap.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
output_path = "output_video.mp4"
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
results = model(img, size=640)
detections = results.pred[0].cpu().numpy()
for *xyxy, conf, cls in detections:
x1, y1, x2, y2 = map(int, xyxy)
class_id = int(cls)
color = colors[class_id].tolist()
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
label = f"{model.names[class_id]} {conf:.2f}"
cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
out.write(frame)
cap.release()
out.release()
return output_path
# Gradio Interface
with gr.Blocks(title="YOLOv5 Object Detection") as demo:
gr.Markdown("# YOLOv5 Object Detection (Image & Video)")
with gr.Tab("Image Detection"):
img_input = gr.Image(label="Upload Image", type="numpy")
img_output = gr.Image(label="Detected Objects", type="numpy")
img_submit = gr.Button("Detect Objects")
img_submit.click(fn=detect_objects, inputs=img_input, outputs=img_output)
with gr.Tab("Video Detection"):
vid_input = gr.Video(label="Upload Video")
vid_output = gr.Video(label="Processed Video")
vid_submit = gr.Button("Process Video")
vid_submit.click(fn=process_video, inputs=vid_input, outputs=vid_output)
demo.launch()
|