Created
November 25, 2021 00:37
-
-
Save cipher982/03a11a4309dd224c4e176e2d25daae03 to your computer and use it in GitHub Desktop.
minimum reproducible Python example for YOLOv5 ONNX inference
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# Setup imports | |
import os | |
from pathlib import Path | |
import cv2 | |
from matplotlib import pyplot as plt | |
import numpy as np | |
import onnx | |
import onnxruntime | |
import torch | |
from utils.augmentations import letterbox | |
from utils.general import non_max_suppression, scale_coords, xywh2xyxy | |
from utils.plots import Annotator, colors | |
%matplotlib inline | |
# Load class names | |
class_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', | |
'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', | |
'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', | |
'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', | |
'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', | |
'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', | |
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', | |
'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', | |
'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', | |
'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] | |
# Load the ONNX model | |
ONNX_PATH = "./yolov5s.onnx" | |
IMG_PATH = "./example_image.jpg" | |
HALF = False | |
# Load a sample image | |
# If you already have an image as an array skip to letterbox step | |
def read_img(fpath, half): | |
img0 = cv2.imread(fpath) | |
img = letterbox(img0, new_shape=(960, 960))[0] | |
img = np.moveaxis(img, -1, 0) | |
img = torch.from_numpy(img).to("cpu") | |
img = img.half() if half else img.float() | |
img /= 255 # 0 - 255 to 0.0 - 1.0 | |
if len(img.shape) == 3: | |
img = img[None] | |
return img.numpy(), img0 | |
img, img0 = read_img(IMG_PATH, HALF) | |
# Run inference with ONNX | |
ort_session = onnxruntime.InferenceSession(ONNX_PATH) | |
y = ort_session.run([ort_session.get_outputs()[0].name], {ort_session.get_inputs()[0].name: img})[0] | |
y = torch.tensor(y) if isinstance(y, np.ndarray) else y | |
# Run ONNX specific modifications to the model output | |
pred = non_max_suppression(y) | |
# Interpret predictions and show image | |
def show_img(pred, im_path, img0, | |
img, names, view_img=True, | |
hide_labels=False, hide_conf=False, | |
): | |
s = f'image {0}/{0} {im_path}: ' | |
for i, det in enumerate(pred): # per image | |
print(det) | |
p, im0, frame = im_path, img0.copy(), img | |
#p = Path(p) # to Path | |
#save_path = str(save_dir / p.name) # im.jpg | |
s += '%gx%g ' % img.shape[2:] # print string | |
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh | |
imc = im0 # for save_crop | |
annotator = Annotator(im0, line_width=3, example=str(names)) | |
if 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 += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string | |
print(s) | |
# Write results | |
for *xyxy, conf, cls in reversed(det): | |
if view_img: # Add bbox to image | |
c = int(cls) # integer class | |
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}') | |
annotator.box_label(xyxy, label, color=colors(c, True)) | |
im0 = annotator.result() | |
plt.imshow(im0) | |
show_img(pred, Path(IMG_PATH), img0, img, class_names) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment