Spaces:
Sleeping
Sleeping
File size: 1,740 Bytes
936b8a6 caff61e 936b8a6 46e3370 a29d5e2 936b8a6 a29d5e2 caff61e 46e3370 a29d5e2 caff61e a29d5e2 46e3370 a29d5e2 46e3370 a29d5e2 46e3370 a29d5e2 46e3370 a29d5e2 46e3370 936b8a6 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 47 |
import os
import torch
import gradio as gr
import cv2
import numpy as np
from PIL import Image
from torchvision.transforms import functional as F
# Clone yolov5 if not present
if not os.path.exists("yolov5"):
os.system("git clone https://github.com/ultralytics/yolov5.git")
os.system("pip install -r yolov5/requirements.txt")
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)
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()
|