Aumkeshchy2003's picture
Update app.py
6ade533 verified
raw
history blame
4.9 kB
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)
# Select device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Use YOLOv5 Nano for speed
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)
# Optimize model for speed
model.conf = 0.3 # Confidence threshold
model.iou = 0.3 # IoU threshold for Non-Maximum Suppression (NMS)
model.classes = None # Detect all classes
model.eval()
if device.type == "cuda":
print("Using FP16 precision for inference (high speed, lower accuracy)")
model.half() # Enable FP16 for faster inference
torch.set_num_threads(os.cpu_count()) # Optimize CPU threading
# Pre-generate colors for bounding boxes
np.random.seed(42)
colors = np.random.uniform(0, 255, size=(len(model.names), 3))
# FPS tracking
total_inference_time = 0
inference_count = 0
def preprocess_image(image):
"""Prepares image for YOLOv5 detection while maintaining aspect ratio."""
h, w, _ = image.shape
scale = 640 / max(h, w)
new_w, new_h = int(w * scale), int(h * scale)
resized_image = cv2.resize(image, (new_w, new_h))
padded_image = np.full((640, 640, 3), 114, dtype=np.uint8) # Gray padding
padded_image[:new_h, :new_w] = resized_image
return cv2.cvtColor(padded_image, cv2.COLOR_RGB2BGR) # Convert to BGR for OpenCV
def detect_objects(image):
global total_inference_time, inference_count
if image is None:
return None
start_time = time.time()
# Preprocess image
image = preprocess_image(image)
with torch.inference_mode(): # Faster than torch.no_grad()
results = model(image, size=640)
inference_time = time.time() - start_time
total_inference_time += inference_time
inference_count += 1
avg_inference_time = total_inference_time / inference_count
detections = results.xyxy[0].cpu().numpy() # Use xyxy format
output_image = image.copy()
for *xyxy, conf, cls in detections:
x1, y1, x2, y2 = map(int, xyxy)
class_id = int(cls)
color = colors[class_id].tolist()
# Keep bounding boxes within image bounds
x1, y1, x2, y2 = max(0, x1), max(0, y1), min(640, x2), min(640, y2)
# Draw bounding box
cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
label = f"{model.names[class_id]} {conf:.2f}"
font_scale, font_thickness = 0.9, 2
(w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
# Label background
cv2.rectangle(output_image, (x1, y1 - h - 10), (x1 + w + 10, y1), color, -1)
cv2.putText(output_image, label, (x1 + 5, y1 - 5),
cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
fps = 1 / inference_time
# Display FPS
overlay = output_image.copy()
cv2.rectangle(overlay, (10, 10), (300, 80), (0, 0, 0), -1)
output_image = cv2.addWeighted(overlay, 0.6, output_image, 0.4, 0)
cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (20, 70),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
return output_image
# Gradio UI
example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
os.makedirs("examples", exist_ok=True)
with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
gr.Markdown("""
# Optimized YOLOv5 Object Detection
Detects objects using YOLOv5 with enhanced visualization and FPS tracking.
""")
with gr.Row():
with gr.Column(scale=1):
input_image = gr.Image(label="Input Image", type="numpy")
submit_button = gr.Button("Submit", variant="primary")
clear_button = gr.Button("Clear")
with gr.Column(scale=1):
output_image = gr.Image(label="Detected Objects", type="numpy")
gr.Examples(
examples=example_images,
inputs=input_image,
outputs=output_image,
fn=detect_objects,
cache_examples=True
)
submit_button.click(fn=detect_objects, inputs=input_image, outputs=output_image)
clear_button.click(lambda: (None, None), None, [input_image, output_image])
demo.launch()