cipher982 · November 25, 2021 00:37
diff --git a/onnx_inference_yolov5.py b/onnx_inference_yolov5.py
 # 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)
	# 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)