python使用onnx模型进行推理

流程描述

1. 定义了一个名为 letterbox 的函数，用于图像缩放和填充。
2. 定义了一个包含80个物体类别名称的列表 names，和一个包含每个类别对应的随机颜色的字典 colors。
3. 将原始图像 img 从 BGR 颜色空间转换为 RGB 颜色空间。
4. 复制一份原始图像，并使用 letterbox 函数对其进行缩放和填充。
5. 将缩放和填充后的图像的维度由 (height, width, channel) 转换为 (channel, height, width)。
6. 在第一维（即 batch 维度）上添加一维，将缩放和填充后的图像转换为大小为 (1, channel, height, width) 的张量。
7. 将图像转换为 numpy 数组，并将其变为连续的存储方式。
8. 将图像中的像素值除以255，将其归一化到 [0,1] 的范围内。
9. 获取模型输出的名称列表，并将其保存在 outname 变量中。
10. 获取模型输入的名称列表，并将其保存在 inname 变量中。
11. 创建一个字典 inp，以模型输入的名称为键，缩放和填充后的图像数组为值。

代码实现，模型使用的是yolov7

引入依赖

import cv2
import time
import requests
import random
import numpy as np
import onnxruntime as ort
from PIL import Image
from pathlib import Path
from collections import OrderedDict,namedtuple

加载模型

cuda = False
w = "./yolov7-tiny.onnx"
img = cv2.imread('../inference/images/horses.jpg')

providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
session = ort.InferenceSession(w, providers=providers)

定义预处理

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32):
    # 获取图像数据的宽度和高度的,[:2]则是取前两个元素，即获取宽度和高度
    shape = im.shape[:2]  
    
    # 判断 new_shape 是否为整数类型。如果是整数类型，则将 new_shape 转换成一个元组，元组中的两个元素均为 new_shape。这样做的目的是为了方便后续的处理，因为在处理图像时，往往需要统一图像的宽度和高度
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # 缩放比例 (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # 在训练时只将图像缩小到指定的大小，而不进行放大。这样做是为了获得更好的验证集均值平均精度（val mAP），因为放大图像会使目标的细节变得模糊，从而影响算法的性能评估。因此，训练时只进行缩小操作，可以更好地评估模型在实际应用中的性能
        r = min(r, 1.0)

    # 在卷积神经网络中，输入图像的大小需要与卷积核大小匹配，因此可能需要在图像周围添加填充来达到相应的大小。填充的大小可以根据具体情况进行计算，并根据需要进行调整。
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # 我们可以在图像的左右两侧或上下两侧添加一些像素点，以使图像大小变为模型输入所需的大小。"wh padding" 中的 "wh" 指代 "width"（宽度）和 "height"（高度）两个方向。

    if auto:  # （最小矩形）指的是能够包含所有目标的最小矩形
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # 计算输入宽度 dw 与步长 stride 的模（余数）
        
    # 在卷积操作中，如果输入图像大小不能被卷积核大小整除，为了能够进行卷积操作，通常需要对输入图像进行 padding 处理，使其能够被卷积核整除。而在进行 padding 处理后，卷积核的扫描位置也会受到影响，需要重新计算。
    # 因此，np.mod(dw, stride) 的作用就是计算在进行 padding 处理后，卷积核在水平方向上能够扫描的位置范围。如果 dw 能够被 stride 整除，则卷积核在水平方向上的扫描位置范围等于 dw。
    # 否则，卷积核在水平方向上的扫描位置范围等于 dw 减去 dw 除以 stride 的余数。

    dw /= 2  # 将填充均匀地分成宽高两个方向
    dh /= 2

    if shape[::-1] != new_unpad:  # 原始图像的宽高与经过缩放和填充后的新尺寸不一致
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, r, (dw, dh)

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']
colors = {name:[random.randint(0, 255) for _ in range(3)] for i,name in enumerate(names)}

# 将BGR（Blue, Green, Red）图像转换为RGB（Red, Green, Blue）
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

image = img.copy()
image, ratio, dwdh = letterbox(image, auto=False)

 #将 image 的维度顺序从 (height, width, channels) 转变为 (channels, height, width)
image = image.transpose((2, 0, 1))

# np.expand_dims(image, 0) 表示在 image 的最前面（索引为 0 的位置）插入一个新维度。
# 例如，如果 image 是一个形状为 (3, 4, 5) 的三维数组，则 np.expand_dims(image, 0) 的结果是一个形状为 (1, 3, 4, 5) 的四维数组。
# 这里插入新维度的目的可能是为了将 image 转换为一个 batch，即一批多个样本的输入，第一个维度通常表示 batch size，因此插入一个维度表示 batch size 为 1
image = np.expand_dims(image, 0)

# 将 image 转换为一个连续的数组
image = np.ascontiguousarray(image)

im = image.astype(np.float32)

# 归一化，将像素值从[0, 255]的范围缩放到[0, 1]的范围
im /= 255

im.shape

# 从session对象的输出中获取每个输出的名称，并将其存储在一个列表中
outname = [i.name for i in session.get_outputs()]


inname = [i.name for i in session.get_inputs()]

inp = {inname[0]:im}

模型推理

outputs = session.run(outname, inp)[0]
outputs 
打印：
array([[ 0.0000000e+00,  3.6190897e+02,  2.8389810e+02,  4.9353049e+02,
         3.9562729e+02,  1.7000000e+01,  9.2383689e-01],
       [ 0.0000000e+00, -1.0339203e+00,  2.6461755e+02,  2.6221344e+02,
         4.4826135e+02,  1.7000000e+01,  9.2106485e-01],
       [ 0.0000000e+00,  2.1546234e+02,  2.7049042e+02,  3.5089423e+02,
         4.1111603e+02,  1.7000000e+01,  7.6384640e-01],
       [ 0.0000000e+00, -9.6609497e-01,  2.6136026e+02,  1.2928017e+02,
         3.3445981e+02,  1.7000000e+01,  6.9170284e-01],
       [ 0.0000000e+00,  3.0596024e+02,  2.8081891e+02,  3.7848898e+02,
         3.7234491e+02,  1.7000000e+01,  4.6638113e-01]], dtype=float32)

画锚框

ori_images = [img.copy()]

for i,(batch_id,x0,y0,x1,y1,cls_id,score) in enumerate(outputs):
    image = ori_images[int(batch_id)]
    box = np.array([x0,y0,x1,y1])
    box -= np.array(dwdh*2)
    box /= ratio
    box = box.round().astype(np.int32).tolist()
    cls_id = int(cls_id)
    score = round(float(score),3)
    name = names[cls_id]
    color = colors[name]
    name += ' '+str(score)
    cv2.rectangle(image,box[:2],box[2:],color,2)
    cv2.putText(image,name,(box[0], box[1] - 2),cv2.FONT_HERSHEY_SIMPLEX,0.75,[225, 255, 255],thickness=2)  

Image.fromarray(ori_images[0])

效果展示

等比缩放后不足640的地方填充指定颜色(114, 114, 114)