Spaces:
Running
Running
File size: 2,034 Bytes
caff61e e82b28e bccf53b dc80d48 4fa263e 248b9ce caff61e f54253a e82b28e 4fa263e 248b9ce 4fa263e e9a3a08 36e1064 4fa263e f54253a a29d5e2 4fa263e e82b28e 4fa263e f54253a 4fa263e 248b9ce 4fa263e e82b28e f54253a 4fa263e f54253a 4fa263e 73df658 4fa263e a29d5e2 46e3370 4fa263e e82b28e 1195707 4fa263e e82b28e 46e3370 f54253a 73df658 |
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 |
import torch
import cv2
import numpy as np
import gradio as gr
from PIL import Image
import random
# Load YOLOv5 model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device)
# Get class names from the model
CLASS_NAMES = model.names
# Generate consistent colors for each class
random.seed(42) # Fix the seed for consistent colors
CLASS_COLORS = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in CLASS_NAMES}
def preprocess_image(image):
"""Convert numpy image to PIL format for YOLOv5 processing."""
image = Image.fromarray(image)
image = image.convert("RGB")
return image
def detect_objects(image):
"""Detect objects in the image and draw bounding boxes with consistent colors."""
image = preprocess_image(image)
results = model([image]) # YOLOv5 inference
image = np.array(image) # Convert PIL image back to numpy for OpenCV
for *box, conf, cls in results.xyxy[0]:
x1, y1, x2, y2 = map(int, box)
class_name = CLASS_NAMES[int(cls)]
confidence = conf.item() * 100
color = CLASS_COLORS[class_name] # Use pre-generated consistent color
# Draw bounding box
cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)
# Display class label with confidence score
label = f"{class_name} ({confidence:.1f}%)"
cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
1, color, 3, cv2.LINE_AA)
return image
# Create Gradio Interface
iface = gr.Interface(
fn=detect_objects,
inputs=gr.Image(type="numpy", label="Upload Image"),
outputs=gr.Image(type="numpy", label="Detected Objects"),
title="Object Detection with YOLOv5",
description="Use webcam or upload an image to detect objects.",
allow_flagging="never",
examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
)
# Launch the app
iface.launch()
|