Spaces:
Running
Running
File size: 1,567 Bytes
46e3370 caff61e 46e3370 a29d5e2 caff61e 46e3370 a29d5e2 caff61e a29d5e2 46e3370 a29d5e2 46e3370 a29d5e2 46e3370 a29d5e2 46e3370 a29d5e2 46e3370 |
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 |
import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
from torchvision.transforms import functional as F
from yolov5.utils.general import non_max_suppression
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device)
model.eval()
def preprocess_image(image):
image = image.convert("RGB")
image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
return image_tensor
def draw_boxes(image, outputs, threshold=0.3):
image = np.array(image)
h, w, _ = image.shape
for box in outputs:
score, label, x1, y1, x2, y2 = box[4].item(), int(box[5].item()), box[0].item(), box[1].item(), box[2].item(), box[3].item()
if score > threshold:
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
text = f"{model.names[label]}: {score:.2f}"
cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
return Image.fromarray(image)
def detect_objects(image):
image_tensor = preprocess_image(image)
outputs = model(image_tensor)
outputs = non_max_suppression(outputs)[0]
return draw_boxes(image, outputs)
iface = gr.Interface(
fn=detect_objects,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
title="YOLO Object Detector",
description="Upload an image to detect objects using YOLOv5."
)
if __name__ == "__main__":
iface.launch()
|