YOLOv6-Openvino和ONNXRuntime推理【CPU】

1 環境：

CPU：i5-12500
Python：3.8.18

2 安裝Openvino和ONNXRuntime

2.1 Openvino簡介

Openvino是由Intel開發的專門用于優化和部署人工智能推理的半開源的工具包，主要用于對深度推理做優化。

Openvino內部集成了Opencv、TensorFlow模塊，除此之外它還具有強大的Plugin開發框架，允許開發者在Openvino之上對推理過程做優化。

Openvino整體框架為：Openvino前端→ Plugin中間層→ Backend后端
Openvino的優點在于它屏蔽了后端接口，提供了統一操作的前端API，開發者可以無需關心后端的實現，例如后端可以是TensorFlow、Keras、ARM-NN，通過Plugin提供給前端接口調用，也就意味著一套代碼在Openvino之上可以運行在多個推理引擎之上，Openvino像是類似聚合一樣的開發包。

2.2 ONNXRuntime簡介

ONNXRuntime是微軟推出的一款推理框架，用戶可以非常便利的用其運行一個onnx模型。ONNXRuntime支持多種運行后端包括CPU，GPU，TensorRT，DML等。可以說ONNXRuntime是對ONNX模型最原生的支持。

雖然大家用ONNX時更多的是作為一個中間表示，從pytorch轉到onnx后直接喂到TensorRT或MNN等各種后端框架，但這并不能否認ONNXRuntime是一款非常優秀的推理框架。而且由于其自身只包含推理功能（最新的ONNXRuntime甚至已經可以訓練），通過閱讀其源碼可以解深度學習框架的一些核心功能原理（op注冊，內存管理，運行邏輯等）
總體來看，整個ONNXRuntime的運行可以分為三個階段，Session構造，模型加載與初始化和運行。和其他所有主流框架相同，ONNXRuntime最常用的語言是python，而實際負責執行框架運行的則是C++。

2.3 安裝

pip install openvino -i  https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnxruntime -i  https://pypi.tuna.tsinghua.edu.cn/simple

3 YOLOv6介紹

YOLOv6詳解
文章地址：https://link.csdn.net/?target=https%3A%2F%2Farxiv.org%2Fpdf%2F2209.02976.pdf
Github：https://github.com/meituan/YOLOv6/tree/main

4 基于Openvino和ONNXRuntime推理

下面代碼整個處理過程主要包括：預處理—>推理—>后處理—>畫圖。
假設圖像resize為640×640，
前處理輸出結果維度：(1, 3, 640, 640)；
推理輸出結果維度：(1, 8400, 85)，其中85表示4個box坐標信息+置信度分數+80個類別概率，8400表示80×80+40×40+20×20，不同于v8與v9采用類別里面最大的概率作為置信度score；
后處理輸出結果維度：(5, 6)，其中第一個5表示圖bus.jpg檢出5個目標，第二個維度6表示(x1, y1, x2, y2, conf, cls)。
注：與YOLOv5輸出維度一致，可通用！！！

4.1 全部代碼

import argparse
import time 
import cv2
import numpy as np
from openvino.runtime import Core  # pip install openvino -i  https://pypi.tuna.tsinghua.edu.cn/simple
import onnxruntime as ort  # 使用onnxruntime推理用上，pip install onnxruntime，默認安裝CPU# COCO默認的80類
CLASSES = ['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']class OpenvinoInference(object):def __init__(self, onnx_path):self.onnx_path = onnx_pathie = Core()self.model_onnx = ie.read_model(model=self.onnx_path)self.compiled_model_onnx = ie.compile_model(model=self.model_onnx, device_name="CPU")self.output_layer_onnx = self.compiled_model_onnx.output(0)def predirts(self, datas):predict_data = self.compiled_model_onnx([datas])[self.output_layer_onnx]return predict_dataclass YOLOv6:"""YOLOv6 object detection model class for handling inference and visualization."""def __init__(self, onnx_model, imgsz=(640, 640), infer_tool='openvino'):"""Initialization.Args:onnx_model (str): Path to the ONNX model."""self.infer_tool = infer_toolif self.infer_tool == 'openvino':# 構建openvino推理引擎self.openvino = OpenvinoInference(onnx_model)self.ndtype = np.singleelse:# 構建onnxruntime推理引擎self.ort_session = ort.InferenceSession(onnx_model,providers=['CUDAExecutionProvider', 'CPUExecutionProvider']if ort.get_device() == 'GPU' else ['CPUExecutionProvider'])# Numpy dtype: support both FP32 and FP16 onnx modelself.ndtype = np.half if self.ort_session.get_inputs()[0].type == 'tensor(float16)' else np.singleself.classes = CLASSES  # 加載模型類別self.model_height, self.model_width = imgsz[0], imgsz[1]  # 圖像resize大小self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))  # 為每個類別生成調色板def __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45):"""The whole pipeline: pre-process -> inference -> post-process.Args:im0 (Numpy.ndarray): original input image.conf_threshold (float): confidence threshold for filtering predictions.iou_threshold (float): iou threshold for NMS.Returns:boxes (List): list of bounding boxes."""# 前處理Pre-processt1 = time.time()im, ratio, (pad_w, pad_h) = self.preprocess(im0)print('預處理時間：{:.3f}s'.format(time.time() - t1))# 推理 inferencet2 = time.time()if self.infer_tool == 'openvino':preds = self.openvino.predirts(im)else:preds = self.ort_session.run(None, {self.ort_session.get_inputs()[0].name: im})[0]print('推理時間：{:.2f}s'.format(time.time() - t2))# 后處理Post-processt3 = time.time()boxes = self.postprocess(preds,im0=im0,ratio=ratio,pad_w=pad_w,pad_h=pad_h,conf_threshold=conf_threshold,iou_threshold=iou_threshold,)print('后處理時間：{:.3f}s'.format(time.time() - t3))return boxes# 前處理，包括：resize, pad, HWC to CHW，BGR to RGB，歸一化，增加維度CHW -> BCHWdef preprocess(self, img):"""Pre-processes the input image.Args:img (Numpy.ndarray): image about to be processed.Returns:img_process (Numpy.ndarray): image preprocessed for inference.ratio (tuple): width, height ratios in letterbox.pad_w (float): width padding in letterbox.pad_h (float): height padding in letterbox."""# Resize and pad input image using letterbox() (Borrowed from Ultralytics)shape = img.shape[:2]  # original image shapenew_shape = (self.model_height, self.model_width)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])ratio = r, rnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # wh paddingif shape[::-1] != new_unpad:  # resizeimg = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))  # 填充# Transforms: HWC to CHW -> BGR to RGB -> div(255) -> contiguous -> add axis(optional)img = np.ascontiguousarray(np.einsum('HWC->CHW', img)[::-1], dtype=self.ndtype) / 255.0img_process = img[None] if len(img.shape) == 3 else imgreturn img_process, ratio, (pad_w, pad_h)# 后處理，包括：閾值過濾與NMSdef postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold):"""Post-process the prediction.Args:preds (Numpy.ndarray): predictions come from ort.session.run().im0 (Numpy.ndarray): [h, w, c] original input image.ratio (tuple): width, height ratios in letterbox.pad_w (float): width padding in letterbox.pad_h (float): height padding in letterbox.conf_threshold (float): conf threshold.iou_threshold (float): iou threshold.Returns:boxes (List): list of bounding boxes."""# (Batch_size, Num_anchors, xywh_score_conf_cls), v5和v6_1.0的[..., 4]是置信度分數，v8v9采用類別里面最大的概率作為置信度scorex = preds  # outputs: predictions (1, 8400, 85)# Predictions filtering by conf-thresholdx = x[x[..., 4] > conf_threshold]# Create a new matrix which merge these(box, score, cls) into one# For more details about `numpy.c_()`: https://numpy.org/doc/1.26/reference/generated/numpy.c_.htmlx = np.c_[x[..., :4], x[..., 4], np.argmax(x[..., 5:], axis=-1)]# NMS filtering# 經過NMS后的值, np.array([[x, y, w, h, conf, cls], ...]), shape=(-1, 4 + 1 + 1)x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]# 重新縮放邊界框，為畫圖做準備if len(x) > 0:# Bounding boxes format change: cxcywh -> xyxyx[..., [0, 1]] -= x[..., [2, 3]] / 2x[..., [2, 3]] += x[..., [0, 1]]# Rescales bounding boxes from model shape(model_height, model_width) to the shape of original imagex[..., :4] -= [pad_w, pad_h, pad_w, pad_h]x[..., :4] /= min(ratio)# Bounding boxes boundary clampx[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])return x[..., :6]  # boxeselse:return []# 繪框def draw_and_visualize(self, im, bboxes, vis=False, save=True):"""Draw and visualize results.Args:im (np.ndarray): original image, shape [h, w, c].bboxes (numpy.ndarray): [n, 6], n is number of bboxes.vis (bool): imshow using OpenCV.save (bool): save image annotated.Returns:None"""# Draw rectangles for (*box, conf, cls_) in bboxes:# draw bbox rectanglecv2.rectangle(im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),self.color_palette[int(cls_)], 1, cv2.LINE_AA)cv2.putText(im, f'{self.classes[int(cls_)]}: {conf:.3f}', (int(box[0]), int(box[1] - 9)),cv2.FONT_HERSHEY_SIMPLEX, 0.7, self.color_palette[int(cls_)], 2, cv2.LINE_AA)# Show imageif vis:cv2.imshow('demo', im)cv2.waitKey(0)cv2.destroyAllWindows()# Save imageif save:cv2.imwrite('demo.jpg', im)if __name__ == '__main__':# Create an argument parser to handle command-line argumentsparser = argparse.ArgumentParser()parser.add_argument('--model', type=str, default='yolov6s.onnx', help='Path to ONNX model')parser.add_argument('--source', type=str, default=str('bus.jpg'), help='Path to input image')parser.add_argument('--imgsz', type=tuple, default=(640, 640), help='Image input size')parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')parser.add_argument('--iou', type=float, default=0.45, help='NMS IoU threshold')parser.add_argument('--infer_tool', type=str, default='openvino', choices=("openvino", "onnxruntime"), help='選擇推理引擎')args = parser.parse_args()# Build modelmodel = YOLOv6(args.model, args.imgsz, args.infer_tool)# Read image by OpenCVimg = cv2.imread(args.source)# Inferenceboxes = model(img, conf_threshold=args.conf, iou_threshold=args.iou)# Visualizeif len(boxes) > 0:model.draw_and_visualize(img, boxes, vis=False, save=True)

4.2 結果

在這里插入圖片描述

具體時間消耗：

預處理時間：0.005s（包含Pad）
推理時間：0.08~0.09s（Openvino）
推理時間：0.11~0.12s（ONNXRuntime）
后處理時間：0.001s
注：640×640下，YOLOv6s.onnx版本為1.0，2.0以上推理結果沒有直接輸出置信度分數，需參考v8v9處理方式對應修改。