Aumkeshchy2003's picture
Update app.py
6fea677 verified
raw
history blame
6.06 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 if it doesn't exist
os.makedirs("models", exist_ok=True)
# Check device availability - Hugging Face Spaces often provides GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load YOLOv5x model with caching for faster startup
model_path = Path("models/yolov5x.pt")
if model_path.exists():
print(f"Loading model from cache: {model_path}")
model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True,
source="local", path=str(model_path)).to(device)
else:
print("Downloading YOLOv5x model and caching...")
model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
# Cache the model for faster startup next time
torch.save(model.state_dict(), model_path)
# Optimization configurations
model.conf = 0.3 # Confidence threshold of 0.3 as specified
model.iou = 0.3 # NMS IoU threshold of 0.3 as specified
model.classes = None # Detect all 80+ COCO classes
# Optimize for GPU if available
if device.type == "cuda":
# Use mixed precision for performance boost
model.half()
else:
# On CPU, optimize operations
torch.set_num_threads(os.cpu_count())
# Set model to evaluation mode for inference
model.eval()
# Assign fixed colors to each class for consistent visualization
np.random.seed(42) # For reproducible colors
colors = np.random.uniform(0, 255, size=(len(model.names), 3))
# Track performance metrics
total_inference_time = 0
inference_count = 0
def detect_objects(image):
"""
Process input image for object detection using YOLOv5
Args:
image: Input image as numpy array
Returns:
output_image: Image with detection results visualized
"""
global total_inference_time, inference_count
if image is None:
return None
start_time = time.time()
# Create a copy for drawing results
output_image = image.copy()
# Fixed input size for optimal processing
input_size = 640
# Perform inference with no gradient calculation
with torch.no_grad():
# Convert image to tensor for faster processing
results = model(image, size=input_size)
# Record inference time (model processing only)
inference_time = time.time() - start_time
total_inference_time += inference_time
inference_count += 1
avg_inference_time = total_inference_time / inference_count
# Extract detections from first (and only) image
detections = results.pred[0].cpu().numpy()
# Draw each detection on the output image
for *xyxy, conf, cls in detections:
# Extract coordinates and convert to integers
x1, y1, x2, y2 = map(int, xyxy)
class_id = int(cls)
# Get color for this class
color = colors[class_id].tolist()
cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 4)
# Create label with class name and confidence score
label = f"{model.names[class_id]} {conf:.2f}"
font_scale = 0.8
font_thickness = 2
(w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
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, (0, 0, 0), font_thickness + 1)
cv2.putText(output_image, label, (x1 + 5, y1 - 5),
cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness)
# Calculate FPS
fps = 1 / inference_time
fps_overlay = output_image.copy()
cv2.rectangle(fps_overlay, (5, 5), (250, 80), (0, 0, 0), -1)
# Apply the overlay with transparency
alpha = 0.7
output_image = cv2.addWeighted(fps_overlay, alpha, output_image, 1 - alpha, 0)
# Display FPS with larger font
cv2.putText(output_image, f"FPS: {fps:.2f}", (10, 35),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (10, 70),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
return output_image
# Define example images - these will be stored in the same directory as this script
example_images = [
"examples/spring_street_after.jpg",
"examples/pexels-hikaique-109919.jpg"
]
# Make sure example directory exists
os.makedirs("examples", exist_ok=True)
# Create Gradio interface - optimized for Hugging Face Spaces
with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
gr.Markdown("""
# Optimized YOLOv5 Object Detection
This system utilizes YOLOv5 to detect 80+ object types from the COCO dataset.
**Performance Features:**
- Processing speed: Optimized for 30+ FPS at 640x640 resolution
- Confidence threshold: 0.3
- IoU threshold: 0.3
Simply upload an image and click Submit to see the detections!
""")
with gr.Row():
with gr.Column(scale=1):
input_image = gr.Image(label="Input Image", type="numpy")
with gr.Row():
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")
# Example gallery
gr.Examples(
examples=example_images,
inputs=input_image,
outputs=output_image,
fn=detect_objects,
cache_examples=True # Cache for faster response
)
# Set up button event handlers
submit_button.click(fn=detect_objects, inputs=input_image, outputs=output_image)
clear_button.click(lambda: None, None, [input_image, output_image])
# Launch for Hugging Face Spaces
demo.launch()