python使用onnx模型进行推理
创始人
2024-06-03 16:52:40
0次

流程描述

  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 /= 2if 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 borderreturn 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 /= 255im.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 /= ratiobox = 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])

效果展示
在这里插入图片描述

相关内容

热门资讯

前端-session、jwt 目录:   (1)session (2&#x...
linux入门---制作进度条 了解缓冲区 我们首先来看看下面的操作: 我们首先创建了一个文件并在这个文件里面添加了...
关于测试,我发现了哪些新大陆 关于测试 平常也只是听说过一些关于测试的术语,但并没有使用过测试工具。偶然看到编程老师...
前缀和与对数器与二分法 1. 前缀和 假设有一个数组,我们想大量频繁的去访问L到R这个区间的和,...
nodejs:本地安装nvm实... 一、背景-使用不同版本node的原因 vue3+ts、nuxt3版本,node...
JAVA集合知识整理 Java集合知识整理 HashMap相关 HashMap的底层数据结构:jdk1.8之...
无刷直流电机介绍及单片机控制实... 无刷直流电机介绍及单片机控制实例前言基本概念优势与劣势使用寿命基本结构使用单片机控制实例电子调速器&...
fwdiary(2) dp2 1.传纸条  AcWing 275. 传纸条 - AcWing 走两条路,走一条最大的...
常用的DOS命令 常用的DOS命令 DOS(Disk Operating System,磁...
<C++> 类和对象(下) 1.const成员函数将const修饰的“成员函数”称之为const成员函数,cons...