Spaces:
Running
Running
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() | |