pytorch下yolov8打包onnx模型并使用NMS对数据清洗（最后有完整代码）

已于 2024-07-15 15:27:47 修改
阅读量467 收藏 9
点赞数 3
文章标签： 目标检测 YOLO pytorch 人工智能 深度学习
于 2024-07-15 15:21:21 首次发布
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/qq_64809150/article/details/140436333
版权

本文对该模型进onnx打包并在pytorch环境下运行（其实我打包onnx的目的是为了部署，大家可以不打包的）

目录

一、打包onnx

二、定义函数方法

1、数据结构介绍

2、函数方法解释

3、效果图

三、完整代码

一、打包onnx

首先对训练得到的模型进行打包，我用官方训练方法得到的模型默认名是best.pt：

from ultralytics import YOLO
model = YOLO('runs/detect/train/weights/best.pt')
success = model.export(format="onnx", simplify=True)
assert success
print("转换成功")

然后得到一个best.onnx的模型文件，接下来对该模型进行运行处理

二、定义函数方法

定义三个函数方法，分别计算标准输出、非极大值抑制（NMS）、交并比（IuO）

def std_output(pred):
    pred = np.squeeze(pred)
    pred = np.transpose(pred, (1, 0))
    pred_class = pred[..., 4:]
    pred_conf = np.max(pred_class, axis=-1)
    pred = np.insert(pred, 4, pred_conf, axis=-1)
    return pred


def nms(detections, iou_threshold):
    if len(detections) == 0:
        return []

    # 按置信度排序
    detections = detections[detections[:, 4].argsort()[::-1]]

    keep = []
    while len(detections) > 0:
        # 取置信度最高的框
        highest = detections[0]
        keep.append(highest)

        if len(detections) == 1:
            break

        # 计算剩余框与最高置信度框的IOU
        rest_boxes = detections[1:]
        iou = compute_iou(highest, rest_boxes)

        # 保留IOU小于阈值的框
        detections = rest_boxes[iou < iou_threshold]

    return np.array(keep)


def compute_iou(box1, boxes):
    x1 = np.maximum(box1[0] - box1[2] / 2, boxes[:, 0] - boxes[:, 2] / 2)
    y1 = np.maximum(box1[1] - box1[3] / 2, boxes[:, 1] - boxes[:, 3] / 2)
    x2 = np.minimum(box1[0] + box1[2] / 2, boxes[:, 0] + boxes[:, 2] / 2)
    y2 = np.minimum(box1[1] + box1[3] / 2, boxes[:, 1] + boxes[:, 3] / 2)

    inter_area = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    box1_area = box1[2] * box1[3]
    boxes_area = boxes[:, 2] * boxes[:, 3]

    iou = inter_area / (box1_area + boxes_area - inter_area)
    return iou
1、数据结构介绍

这部分就要对多维数组结构进行说明，比如模型输出数组形状是(1, 7, 8400)，这是yolov8的默认输出，那么可以看作有1个7*8400的矩阵，只需要看7*8400这个矩阵，其中8400是锚框数量，7是包括锚框4个坐标+3个识别目标的类型（可以看我上一篇文章配置文件是不是有3个类型），运行模型得到的数组前4位是锚框在图上的四点坐标，后面的位数是置信度，也就是告诉我们这个锚框里面是什么的可能性最大，有多少种类的标签，那么后面就有多少种类的置信度，所以这个数据长度是根据你的模型能识别多少标签来决定的，但是至少有4位是坐标。

2、函数方法解释

接下来对三个方法进行解释，首先是std_output，如上述所言，我们需要去除第一个维度然后进行转置操作，为什么转置？因为8400*7的矩阵便于切片操作，这样每一行都是一个锚框，否则处理起来异常麻烦。切片后计算置信度，我们只需要4位坐标后面的3个置信度，接着计算每个预测框中所有类别的最大置信度，算完后把最大置信度重新插入原来的数组中，并且插入位置是在4个坐标后面，所以该函数处理后输出的数组形状是(8400, 8)

接下来是NMS和IoU，先解释原理，假设图上有一只小狗，蓝色的锚框覆盖整个小狗（这是最大框，计算来自刚刚插入的第5位最大置信度），而红色的锚框覆盖小狗的半身，那么这段代码会计算这两个锚框的交并比，即交集面积除以并集面积，如果这个交并比小于我设定的比例（比如本案例中我的交并比设置为0.2），那么这个红色锚框就会被保留下来并进行下一轮计算，而最大值的那个锚框会加入keep数组中。8400个锚框最终都会被筛选剩下置信度最大的那个，其实这个算法就是选取局部最大值的思想。在上述的NMS方法中的循环函数就是这样和我的IoU进行比较，大的锚框保留下来，而compute_iou方法就是返回一个包含每个框与最大置信度框的 IoU 的数组，提供给NMS方法进行比较。

最后是上述函数方法的使用：

# 加载ONNX模型
onnx_model_path = 'runs/detect/train4/weights/best.onnx'
ort_session = onnxruntime.InferenceSession(onnx_model_path)

# 加载并预处理图像
image_path = 'ricedata/test/images/077.jpg'
image = Image.open(image_path)
preprocess = transforms.Compose([
    transforms.Resize((640, 640)),
    transforms.ToTensor(),
])
input_data = preprocess(image).unsqueeze(0).numpy()  # 添加批处理维度并转换为NumPy数组

# 进行推理
outputs = ort_session.run(None, {'images': input_data})  # 更新输入名称为模型期望的名称

# 获取第一个输出数组
output_array = outputs[0]

# 解析检测结果
num_detections = output_array.shape[2]

# 定义置信度阈值
confidence_threshold = 0.1
new_array = std_output(output_array)

# 过滤低置信度的检测框
filtered_detections = new_array[new_array[:, 4] > confidence_threshold]

# 进行非极大值抑制
nms_detections = nms(filtered_detections, iou_threshold=0.2)
3、效果图

最后在图上画出处理后的锚框就可以了，效果图如下（锚框颜色浅了一点哈哈哈）：

三、完整代码

# from ultralytics import YOLO
# model = YOLO('runs/detect/train3/weights/best.pt')
# success = model.export(format="onnx", simplify=True)
# assert success
# print("转换成功")

import onnxruntime
import numpy as np
from PIL import Image, ImageDraw
import torchvision.transforms as transforms


def std_output(pred):
    pred = np.squeeze(pred)
    pred = np.transpose(pred, (1, 0))
    pred_class = pred[..., 4:]
    pred_conf = np.max(pred_class, axis=-1)
    pred = np.insert(pred, 4, pred_conf, axis=-1)
    return pred


def nms(detections, iou_threshold):
    if len(detections) == 0:
        return []

    # 按置信度排序
    detections = detections[detections[:, 4].argsort()[::-1]]

    keep = []
    while len(detections) > 0:
        # 取置信度最高的框
        highest = detections[0]
        keep.append(highest)

        if len(detections) == 1:
            break

        # 计算剩余框与最高置信度框的IOU
        rest_boxes = detections[1:]
        iou = compute_iou(highest, rest_boxes)

        # 保留IOU小于阈值的框
        detections = rest_boxes[iou < iou_threshold]
    return np.array(keep)


def compute_iou(box1, boxes):
    x1 = np.maximum(box1[0] - box1[2] / 2, boxes[:, 0] - boxes[:, 2] / 2)
    y1 = np.maximum(box1[1] - box1[3] / 2, boxes[:, 1] - boxes[:, 3] / 2)
    x2 = np.minimum(box1[0] + box1[2] / 2, boxes[:, 0] + boxes[:, 2] / 2)
    y2 = np.minimum(box1[1] + box1[3] / 2, boxes[:, 1] + boxes[:, 3] / 2)

    inter_area = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
    box1_area = box1[2] * box1[3]
    boxes_area = boxes[:, 2] * boxes[:, 3]

    iou = inter_area / (box1_area + boxes_area - inter_area)
    return iou


# 加载ONNX模型
onnx_model_path = 'runs/detect/train4/weights/best.onnx'
ort_session = onnxruntime.InferenceSession(onnx_model_path)

# 加载并预处理图像
image_path = 'ricedata/test/images/077.jpg'
image = Image.open(image_path)
preprocess = transforms.Compose([
    transforms.Resize((640, 640)),
    transforms.ToTensor(),
])
input_data = preprocess(image).unsqueeze(0).numpy()  # 添加批处理维度并转换为NumPy数组

# 进行推理
outputs = ort_session.run(None, {'images': input_data})  # 更新输入名称为模型期望的名称

# 获取第一个输出数组
output_array = outputs[0]

# 解析检测结果
num_detections = output_array.shape[2]

# 定义置信度阈值
confidence_threshold = 0.1
new_array = std_output(output_array)

# 过滤低置信度的检测框
filtered_detections = new_array[new_array[:, 4] > confidence_threshold]

# 进行非极大值抑制
nms_detections = nms(filtered_detections, iou_threshold=0.2)

# 打开原始图像
image = Image.open(image_path)

# 创建绘图对象
draw = ImageDraw.Draw(image)

# 计算特征图与原始图像的尺寸比例
original_image_size = image.size
feature_map_size = (640, 640)  
scale_x = original_image_size[0] / feature_map_size[0]
scale_y = original_image_size[1] / feature_map_size[1]

# 绘制经过NMS处理后的检测框
for detection in nms_detections:
    bbox = detection[:4]
    x0 = (bbox[0] - bbox[2] / 2) * scale_x
    y0 = (bbox[1] - bbox[3] / 2) * scale_y
    x1 = (bbox[0] + bbox[2] / 2) * scale_x
    y1 = (bbox[1] + bbox[3] / 2) * scale_y

    # 绘制边界框
    draw.rectangle([(x0, y0), (x1, y1)], outline="red")

# 显示图像
image.show()

上一篇文章： 

pytorch下yolov8模型实现目标检测（全网最简洁快速，一眼懂）-CSDN博客icon-default.png?t=N7T8https://blog.csdn.net/qq_64809150/article/details/140435678?spm=1001.2014.3001.5501

焚詩作薪
关注
  • 3
    点赞
  • 踩
  • 9
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论
pytorch 42 C#使用onnxruntime部署内置nms的yolov8模型
a486259的博客
08-19 1540
在进行目标检测部署时，通常需要自行编码实现对模型预测结果的解码及与预测结果的nms操作。所幸现在的各种部署框架对算子的支持更为灵活，可以在模型内实现预测结果的解码，但仍然需要自行编码实现对预测结果的nms操作。其实在onnx opset===11版本以后，其已支持将nms操作嵌入到模型中，具体可以参考https://hpg123.blog.csdn.net/article/details/131585808，修改代码实现yolov8模型onnx下的无nms导出。
pytorch 41 yolov8的无nsm后处理的onnxruntime部署方案
a486259的博客
07-06 1816
自onnx11版本及以后，onnx支持将nms操作嵌入到模型中。本博文将yolov8模型的nms操作嵌入到模型中可以大幅度的简化模型部署代码，提升模型推理速度（在以往大家都是使用opencv的库进行nms操作，CPU效率必然是不如GPU的）。同时修改了yolov8输出的格式，移除了每一个类别的conf信息，只保留了maxClassConf和maxClassId。后续会补充C++、C#部署案例，目前仅有python onnxruntime部署案例。
OpenCV 调用自定义训练的 YOLO-V8 Onnx 模型
最新发布
小毕超博客
06-30 1788
Onnx 格式是一种开放格式，用于表示深度学习模型。目的是为了促进不同深度学习框架和工具之间的模型互操作性。通过Onnx，开发人员可以更容易地在不同的深度学习框架之间转换模型。本文首先将自定义的 YOLO-V8 Pytorch 模型转为 Onnx 模型，然后通过 OpenCV 调用 Onnx 模型，进行运算推理。
YOLOV8权重onnx文件python部署篇
yxk666的博客
09-06 1688
仅供学习参考使用！！！
yolov8转onnx
qq_39490370的博客
05-16 1172
opencv 用c++加载onnx模型
yolov8：pt 转 onnx
weixin_45963712的博客
01-31 3022
我使用的是虚拟环境，yolov8包已经下载到本地了，因此直接在anaconda prompt 命令行继续安装。首先激活自己的虚拟环境，然后执行安装命令。
三天从YOLOV8关键点检测入门到实战（第三天）——用onnx部署yolov8
chenyuan1111的博客
10-11 2443
其实我们可以看到应用onnx来预测确实是准确性有所降低，这里也是用的最小的模型yolov8n-pose的模型，可能用大的模型会更好。但重点是一些套路，比如onnx推理引擎的使用，opencv可视化的模版，这些套路对以后开发是很有帮助的。
TensorRT-使用TensorRT部署Pytorch版YOLOv4-通过Pytorch2onnx导出onnx模型转TRT部署
02-27
在本文中，我们将深入探讨如何使用TensorRT进行PyTorch版本的YOLOv4模型的部署，通过PyTorch到ONNX的转换，以及进一步将ONNX模型优化为TensorRT执行引擎。这是一个高效的算法部署实践，对于提升目标检测速度和性能...
pytorch导出yolov5 onnx模型用vs2019 c++推理源代码
07-07
本代码是根据yolov5导出的onnx模型，通过c++部署； 本代码是可运行代码，但是项目目录下的包含目录需要...如果调试过程中有其他疑问，可以先看看我的这篇博客：pytorch导出yolov5 onnx模型用vs2019 c++推理保姆级教程
YOLOV8模型转换-ONNX-RKNN
10-12
YOLOV8模型转换至ONNX-RKNN的流程是一个典型的深度学习模型部署步骤，尤其在嵌入式设备如RK3588上运行时，这样的优化和转换至关重要。YOLO（You Only Look Once）系列是实时目标检测系统，而YOLOV8作为其最新版本，...
yolov8的onnx模型转换成tensorRT的engine模型
01-29
Yolov8的ONNX模型是该模型在ONNX格式下的表示，包含了模型的结构和权重信息。 接下来，我们探讨TensorRT。TensorRT通过多种技术如动态形状、INT8量化、层融合等来优化模型，以实现GPU上的高速推理。将ONNX模型转换...
yolov5-onnx-tensorrt:如何使用 TensorRT 运行 yolov5 模型
05-29
将 pytorch 转换为 onnx 和 tensorrt yolov5 模型以在 Jetson AGX Xavier 上运行。 支持推断图像。 支持同时推断多幅图像。 要求 请使用torch>=1.6.0 + onnx==1.8.0 + TensorRT 7.0.0.11 运行代码 代码结构 ...
YOLOv8目标检测：使用ONNX模型进行推理
深度学习领域
02-22 7183
通过以上步骤，我们展示了如何使用YOLOv8进行目标检测的完整流程，从环境配置到代码实现和结果展示。此过程适用于YOLOv8目标检测任意模型进行检测任务。希望这篇博客能够帮助您理解和实现基于YOLOv8的目标检测项目。如果有任何问题或需要进一步的帮助，请随时留言讨论。
YOLOV8模型测试、训练及导出onnx模型
热门推荐
qq_42178122的博客
09-12 1万+
yolov8
YOLOv8 ONNX推理代码讲解——基于numpy实现
swpucwf的博客
05-17 582
本文将详细讲解YOLOv8 ONNX推理的Python代码，包含如何进行本地摄像头、图像和视频推理的实现。通过以上代码，我们可以实现YOLOv8模型在本地进行图像、视频和摄像头的推理。将输入图像按比例缩放到目标尺寸，同时填充背景使得图像保持原比例不变。对图像进行预处理，包括调整大小、归一化以及格式转换。将检测框从中心坐标形式转换为边框坐标形式。将检测框坐标从缩放后图像转换回原始图像。执行非极大值抑制，去除重叠的检测框。将检测框限制在图像范围内。在图像上绘制检测框及标签。将检测结果绘制到图像上。
YOLOV8项目部署
weixin_72608251的博客
03-25 3021
自定义数据集我们在项目下创建一个datasets文件夹数据结构如下所示：我们按照上面创建目录结构，将对应的数据放进去我在这里只是演示，所以放的少修改数据集的路径我们复制一下这个文件，制作一个我们的配置文件修改类别数我们像上次一样复制这个文件，将后面的数字修改为自己的类别数OK，至此数据集制作完毕。
【yolov8部署实战】VS2019环境下使用Onnxruntime环境部署yolov8目标检测|含源码
Yaoyao2024的博客
03-02 3088
部署yolo项目，是我这几个月以来做的事情，最近打算把这几个月试过的方法，踩过的坑，以博客的形式，分享一下。关于下面动态中讲到的如何用opencv部署，我在上一篇博客中已经详细讲到了：【yolov8部署实战】VS2019环境下使用C++和OpenCV环境部署yolo项目|含详细注释源码。这篇博客主要讲讲使用onnxruntime部署 主要参考：https://github.com/Amyheart/yolov5v8-dnn-onnxruntime/tree/main 如果是只想要onnxruntime部署
YOLOv8目标检测实战：ONNX模型转换及TensorRT部署
bai666ai的博客
02-23 622
本课程讲述如何对YOLOv8目标检测的PyTorch权重文件转成ONNX格式并进行TensorRT加速部署。
yolov8导出onnx pytorch
04-28
要将 YOLOv8 模型导出为 ONNX 格式，需要完成以下步骤： 1. 首先，确保已经安装好 PyTorch 和 ONNX 工具包。 2. 加载 YOLOv8 模型，可以使用 PyTorch 提供的 torch.load() 函数加载模型。 3. 将模型转换为 ONNX 格式，可以使用 PyTorch 提供的 torch.onnx.export() 函数。在函数中，需要指定模型、输入张量、输出张量、导出路径等参数。 以下是示例代码： ```python import torch import torchvision import onnx # 加载 YOLOv8 模型 model = torch.load('yolov8.pth') # 定义输入和输出张量 input_tensor = torch.randn(1, 3, 416, 416) output_tensor = model(input_tensor) # 将模型转换为 ONNX 格式 torch.onnx.export(model, input_tensor, 'yolov8.onnx', verbose=True) ``` 在上述代码中，我们首先加载了 YOLOv8 模型，然后定义了输入和输出张量。最后，使用 torch.onnx.export() 函数将模型导出为 ONNX 格式，并指定导出路径为 'yolov8.onnx'。
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

焚詩作薪

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付：¥1
获取中
扫码支付

您的余额不足，请更换扫码支付或充值

打赏作者

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值