File size: 3,367 Bytes
da9c9dd d00ffff da9c9dd |
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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import gradio as gr
import os
os.system("wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt")
os.system("wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s6.pt")
from models.experimental import attempt_load
from utils.augmentations import letterbox
from utils.plots import Annotator
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import *
import sys
import numpy as np
import random
def detect(img, weights):
gpu_id="cpu"
device = select_device(device=gpu_id)
model = attempt_load(weights+'.pt', device=device)
torch.no_grad()
model.to(device).eval()
half = False # half precision only supported on CUDA
if half:
model.half()
img_size = 640
# Get names and colors
names = model.names if hasattr(model, 'names') else model.modules.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
if img is None:
sys.exit(0)
# Run inference
t0 = time_sync()
im0 = img.copy()
img = letterbox(img, img_size, stride=int(model.stride.max()), auto=False and True)[0]
img = np.stack(img, 0)
img = img.transpose((2, 0, 1))[::-1] # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device)
if half:
img = img.half()
else:
img = img.float() # if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if len(img.shape) == 3:
img = img[None] # expand for batch dim
# Inference
t1 = time_sync()
pred = model(img, augment=False, profile=False)[0]
# to float
if half:
pred = pred.float()
# Apply NMS
pred = non_max_suppression(
pred, 0.1, 0.5, classes=None, agnostic=False)
t2 = time_sync()
annotator = Annotator(im0, line_width=3, example=str(names))
# Process detections
for i, det in enumerate(pred): # detections per image
s = ''
s += '%gx%g ' % img.shape[2:] # print string
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# show results
for *xyxy, conf, cls in det:
label = '%s %.2f' % (names[int(cls)], conf)
annotator.box_label(xyxy, label, color=colors[int(cls)])
im0 = annotator.result()
# Print time (inference + NMS)
infer_time = t2 - t1
print('%sDone. %s' %
(s, infer_time))
print('Done. (%.3fs)' % (time.time() - t0))
return im0
if __name__ == '__main__':
gr.Interface(detect,[gr.Image(type="numpy"),gr.Dropdown(choices=["yolov5s","yolov5s6"])],
gr.Image(type="numpy"),title="Yolov5",examples=[["data/images/bus.jpg", "yolov5s"]],
description="Gradio based demo for <a href='https://github.com/ultralytics/yolov5' style='text-decoration: underline' target='_blank'>ultralytics/yolov5</a>, new state-of-the-art for real-time object detection").launch()
|