-
YOLOv8-Cls推理详解及部署实现
前言
梳理下 YOLOv8-Cls 的预处理流程，顺便让 tensorRT_Pro 支持 YOLOv8-Cls
一、YOLOv8-Cls推理(Python)
1. YOLOv8-Cls预测
我们先尝试利用官方预训练权重来推理一张图片，看能否成功
在 YOLOv8 主目录下新建 predict-cls.py 预测文件，其内容如下：
import cv2 from ultralytics import YOLO if __name__ == "__main__": model = YOLO("yolov8s-cls.pt") img = cv2.imread("ultralytics/assets/bus.jpg") result = model(img)[0] names = result.names top1_label = result.probs.top1 top5_label = result.probs.top5 top1_conf = result.probs.top1conf top5_conf = result.probs.top5conf top1_name = names[top1_label] print(f"The model predicted category is {top1_name}, label = {top1_label}, confidence = {top1_conf:.4f}")- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
在上述代码中我们通过 opencv 读取了一张图像，并送入模型中推理得到 results，results 中保存着不同任务的结果，我们这里是分类任务，因此只需要拿到对应 1000 个类别中最高置信度的类别标签即可。
模型推理的结果如下所示：
2. YOLOv8-Cls预处理
模型预测成功后我们就需要自己动手来写下 YOLOv8-Cls 的预处理，方便后续在 C++ 上的实现
经过我们的调试分析可知 YOLOv8-Cls 的预处理过程在 ultralytics/data/augment.py 文件中，可以参考：augment.py#L1059
class CenterCrop: """YOLOv8 CenterCrop class for image preprocessing, designed to be part of a transformation pipeline, e.g., T.Compose([CenterCrop(size), ToTensor()]). """ def __init__(self, size=640): """Converts an image from numpy array to PyTorch tensor.""" super().__init__() self.h, self.w = (size, size) if isinstance(size, int) else size def __call__(self, im): """ Resizes and crops the center of the image using a letterbox method. Args: im (numpy.ndarray): The input image as a numpy array of shape HWC. Returns: (numpy.ndarray): The center-cropped and resized image as a numpy array. """ imh, imw = im.shape[:2] m = min(imh, imw) # min dimension top, left = (imh - m) // 2, (imw - m) // 2 return cv2.resize(im[top:top + m, left:left + m], (self.w, self.h), interpolation=cv2.INTER_LINEAR) class ToTensor: """YOLOv8 ToTensor class for image preprocessing, i.e., T.Compose([LetterBox(size), ToTensor()]).""" def __init__(self, half=False): """Initialize YOLOv8 ToTensor object with optional half-precision support.""" super().__init__() self.half = half def __call__(self, im): """ Transforms an image from a numpy array to a PyTorch tensor, applying optional half-precision and normalization. Args: im (numpy.ndarray): Input image as a numpy array with shape (H, W, C) in BGR order. Returns: (torch.Tensor): The transformed image as a PyTorch tensor in float32 or float16, normalized to [0, 1]. """ im = np.ascontiguousarray(im.transpose((2, 0, 1))[::-1]) # HWC to CHW -> BGR to RGB -> contiguous im = torch.from_numpy(im) # to torch im = im.half() if self.half else im.float() # uint8 to fp16/32 im /= 255.0 # 0-255 to 0.0-1.0 return im- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
- 37
- 38
- 39
- 40
- 41
- 42
- 43
- 44
- 45
- 46
- 47
- 48
- 49
它包含如下步骤：
- im[top:top +m, left:left + m]：中心裁剪
- cv2.resize：缩放到 224x224
- transpose(2, 0, 1)[::-1]：HWC → CHW，BGR → RGB
- torch.from_numpy：to Tensor
- im /= 255.0：除以 255.0，归一化
因此我们不难写出对应的预处理代码，如下所示：
def preprocess(img, dst_width=224, dst_height=224): imh, imw = img.shape[:2] m = min(imh, imw) top, left = (imh - m) // 2, (imw - m) // 2 img_pre = img[top:top+m, left:left+m] img_pre = cv2.resize(img_pre, (dst_width, dst_height), interpolation=cv2.INTER_LINEAR) img_pre = (img_pre[...,::-1] / 255.0).astype(np.float32) img_pre = img_pre.transpose(2, 0, 1)[None] img_pre = torch.from_numpy(img_pre) return img_pre- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
经过中心裁剪并 resize 后的图片如下所示：
3. YOLOv8-Cls推理
由于我们经过 softmax 后直接得到的是每个类别的概率值，因此没有后处理一说，YOLOv8-Cls 的推理包括图像预处理、模型推理，其中预处理主要是 中心裁剪和缩放。
完整的推理代码如下：
import cv2 import torch import numpy as np from ultralytics.nn.autobackend import AutoBackend def preprocess(img, dst_width=224, dst_height=224): imh, imw = img.shape[:2] m = min(imh, imw) top, left = (imh - m) // 2, (imw - m) // 2 img_pre = img[top:top+m, left:left+m] img_pre = cv2.resize(img_pre, (dst_width, dst_height), interpolation=cv2.INTER_LINEAR) img_pre = (img_pre[...,::-1] / 255.0).astype(np.float32) img_pre = img_pre.transpose(2, 0, 1)[None] img_pre = torch.from_numpy(img_pre) return img_pre if __name__ == "__main__": img = cv2.imread("ultralytics/assets/bus.jpg") img_pre = preprocess(img) model = AutoBackend(weights="yolov8s-cls.pt") names = model.names probs = model(img_pre)[0] top1_label = int(probs.argmax()) top5_label = (-probs).argsort(0)[:5].tolist() top1_conf = probs[top1_label] top5_conf = probs[top5_label] top1name = names[top1_label] print(f"The model predicted category is {top1name}, label = {top1_label}, confidence = {top1_conf:.4f}")- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
- 37
推理结果如下所示：
至此，我们在 Python 上面完成了 YOLOv8-Cls 的整个推理过程，下面我们去 C++ 上实现。
二、YOLOv8-Cls推理(C++)
C++ 上的实现我们使用的 repo 依旧是 tensorRT_Pro，现在我们就基于 tensorRT_Pro 完成 YOLOv8-Cls 在 C++ 上的推理。
1. ONNX导出
首先我们需要将 YOLOv8-Cls 模型导出为 ONNX，为了适配 tensorRT_Pro 我们需要做一些修改，主要有以下几点：
- 修改输出节点名 output
- 输入输出只让 batch 维度动态，宽高不动态
具体修改如下：
1. 在 ultralytics/engine/exporter.py 文件中改动一处
- 323 行：输出节点名修改为 output
- 326 行：输入只让 batch 维度动态，宽高不动态
- 331 行：输出只让 batch 维度动态，宽高不动态
# ========== exporter.py ========== # ultralytics/engine/exporter.py第323行 # output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0'] # dynamic = self.args.dynamic # if dynamic: # dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640) # if isinstance(self.model, SegmentationModel): # dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400) # dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) # elif isinstance(self.model, DetectionModel): # dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400) # 修改为： output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output'] dynamic = self.args.dynamic if dynamic: dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640) dynamic['output'] = {0: 'batch'} if isinstance(self.model, SegmentationModel): dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400) dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) elif isinstance(self.model, DetectionModel): dynamic['output'] = {0: 'batch'} # shape(1, 84, 8400)- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
以上就是为了适配 tensorRT_Pro 而做出的代码修改，修改好以后，将预训练权重 yolov8-cls.pt 放在 ultralytics-main 主目录下，新建导出文件 export.py，内容如下：
from ultralytics import YOLO model = YOLO("yolov8s-cls.pt") success = model.export(format="onnx", dynamic=True, simplify=True)- 1
- 2
- 3
- 4
- 5
在终端执行如下指令即可完成 onnx 导出：
python export.py- 1
导出过程如下图所示：
可以看到导出的 pytorch 模型的输入 shape 是 1x3x224x224，输出 shape 是 1x1000，符合我们的预期。
导出成功后会在当前目录下生成 yolov8s-cls.onnx 模型，我们可以使用 Netron 可视化工具查看，如下图所示：
可以看到输入节点名是 images，维度是 batchx3x224x224，保证只有 batch 维度动态，输出节点名是 output，维度是 batchx1000，保证只有 batch 维度动态，符合 tensorRT_Pro 的格式。
2. YOLOv8-Cls预处理
之前有提到过 YOLOv8-Cls 的预处理部分主要是中心裁剪加缩放，而在 tensorRT_Pro 中有提供 resize 的实现，我们只需要添加中心裁剪即可。
因此我们不难写出 YOLOv8-Cls 的预处理代码，如下所示：
__global__ void crop_resize_bilinear_and_normalize_kernel( uint8_t* src, int src_line_size, int src_width, int src_height, float* dst, int dst_width, int dst_height, int crop_x, int crop_y, float sx, float sy, Norm norm, int edge ){ int position = blockDim.x * blockIdx.x + threadIdx.x; if (position >= edge) return; int dx = position % dst_width; int dy = position / dst_width; float src_x = (dx + 0.5f) * sx - 0.5f + crop_x; float src_y = (dy + 0.5f) * sy - 0.5f + crop_y; float c0, c1, c2; int y_low = floorf(src_y); int x_low = floorf(src_x); int y_high = limit(y_low + 1, 0, src_height - 1); int x_high = limit(x_low + 1, 0, src_width - 1); y_low = limit(y_low, 0, src_height - 1); x_low = limit(x_low, 0, src_width - 1); int ly = rint((src_y - y_low) * INTER_RESIZE_COEF_SCALE); int lx = rint((src_x - x_low) * INTER_RESIZE_COEF_SCALE); int hy = INTER_RESIZE_COEF_SCALE - ly; int hx = INTER_RESIZE_COEF_SCALE - lx; int w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; float* pdst = dst + dy * dst_width + dx * 3; uint8_t* v1 = src + y_low * src_line_size + x_low * 3; uint8_t* v2 = src + y_low * src_line_size + x_high * 3; uint8_t* v3 = src + y_high * src_line_size + x_low * 3; uint8_t* v4 = src + y_high * src_line_size + x_high * 3; c0 = resize_cast(w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]); c1 = resize_cast(w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]); c2 = resize_cast(w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]); if(norm.channel_type == ChannelType::Invert){ float t = c2; c2 = c0; c0 = t; } if(norm.type == NormType::MeanStd){ c0 = (c0 * norm.alpha - norm.mean[0]) / norm.std[0]; c1 = (c1 * norm.alpha - norm.mean[1]) / norm.std[1]; c2 = (c2 * norm.alpha - norm.mean[2]) / norm.std[2]; }else if(norm.type == NormType::AlphaBeta){ c0 = c0 * norm.alpha + norm.beta; c1 = c1 * norm.alpha + norm.beta; c2 = c2 * norm.alpha + norm.beta; } int area = dst_width * dst_height; float* pdst_c0 = dst + dy * dst_width + dx; float* pdst_c1 = pdst_c0 + area; float* pdst_c2 = pdst_c1 + area; *pdst_c0 = c0; *pdst_c1 = c1; *pdst_c2 = c2; }- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
- 37
- 38
- 39
- 40
- 41
- 42
- 43
- 44
- 45
- 46
- 47
- 48
- 49
- 50
- 51
- 52
- 53
- 54
- 55
- 56
- 57
- 58
相比于 resize 的实现就多了一个偏移，主要是为了做中心裁剪，具体代码可以参考：preprocess_kernel.cu#L49
3. YOLOv8-Cls推理
通过上面对 YOLOv8-Cls 的预处理分析之后，整个推理过程就显而易见了。C++ 上 YOLOv8-Cls 的预处理部分将 resize 简单修改即可。
我们在终端执行如下指令即可完成推理（注意！完整流程博主会在后续内容介绍，这边只是简单演示）
make yolo_cls- 1
编译图解如下所示：
至此，我们在 C++ 上面完成了 YOLOv8-Cls 的整个推理过程，下面我们将完整的走一遍流程。
三、YOLOv8-Cls部署
博主新建了一个仓库 tensorRT_Pro-YOLOv8，该仓库基于 shouxieai/tensorRT_Pro，并进行了调整以支持 YOLOv8 的各项任务，目前已支持分类、检测、分割、姿态点估计任务。
下面我们就来具体看看如何利用 tensorRT_Pro-YOLOv8 这个 repo 完成 YOLOv8-Cls 的推理。
1. 源码下载
tensorRT_Pro-YOLOv8 的代码可以直接从 GitHub 官网上下载，源码下载地址是 https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8，Linux 下代码克隆指令如下：
git clone https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8.git- 1
也可手动点击下载，点击右上角的
Code按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2023/11/7 日，若有改动请参考最新）2. 环境配置
需要使用的软件环境有 TensorRT、CUDA、cuDNN、OpenCV、Protobuf，所有软件环境的安装可以参考 Ubuntu20.04软件安装大全，这里不再赘述，需要各位看官自行配置好相关环境😄，外网访问较慢，这里提供下博主安装过程中的软件安装包下载链接 Baidu Drive【pwd:yolo】🚀🚀🚀
tensorRT_Pro-YOLOv8 提供 CMakeLists.txt 和 Makefile 两种方式编译，二者选一即可
2.1 配置CMakeLists.txt
主要修改五处
1. 修改第 13 行，修改 OpenCV 路径
set(OpenCV_DIR "/usr/local/include/opencv4/")- 1
2. 修改第 15 行，修改 CUDA 路径
set(CUDA_TOOLKIT_ROOT_DIR "/usr/local/cuda-11.6")- 1
3. 修改第 16 行，修改 cuDNN 路径
set(CUDNN_DIR "/usr/local/cudnn8.4.0.27-cuda11.6")- 1
4. 修改第 17 行，修改 tensorRT 路径
set(TENSORRT_DIR "/opt/TensorRT-8.4.1.5")- 1
5. 修改第 20 行，修改 protobuf 路径
set(PROTOBUF_DIR "/home/jarvis/protobuf")- 1
2.2 配置Makefile
主要修改五处
1. 修改第 4 行，修改 protobuf 路径
lean_protobuf := /home/jarvis/protobuf- 1
2. 修改第 5 行，修改 tensorRT 路径
lean_tensor_rt := /opt/TensorRT-8.4.1.5- 1
3. 修改第 6 行，修改 cuDNN 路径
lean_cudnn := /usr/local/cudnn8.4.0.27-cuda11.6- 1
4. 修改第 7 行，修改 OpenCV 路径
lean_opencv := /usr/local- 1
5. 修改第 8 行，修改 CUDA 路径
lean_cuda := /usr/local/cuda-11.6- 1
3. ONNX导出
导出细节可以查看之前的内容，这边不再赘述。记得将导出的 ONNX 模型放在 tensorRT_Pro-YOLOv8/workspace 文件夹下。
4. 源码修改
如果你想推理自己训练的模型还需要修改下源代码，YOLOv8-Cls 模型的推理代码主要在 app_yolo_cls.cpp 文件中，我们就只需要修改这一个文件中的内容即可，源码修改较简单主要有以下几点：
- 1. app_yolo_cls.cpp 187行，“yolov8s-cls” 修改为你导出的 ONNX 模型名
- 2. app_yolo_cls.cpp 105行，“imagenet.txt” 修改为你自训练分类模型的类别 txt 文件
具体修改示例如下：
test(TRT::Model::FP32, "best") // 修改1 187行"yolov8s-cls"改成"best" auto labels = iLogger::split_string(iLogger::load_text_file("custom.txt"), "\n"); // 修改2 105行修改检测类别，为自训练模型的类别名称- 1
- 2
- 3
OK！源码修改好了，Makefile 编译文件也搞定了，ONNX 模型也准备好了，现在可以编译运行了，直接在终端执行如下指令即可：
make yolo_cls- 1
编译过程如下所示：
编译运行成功后在 workspace 文件夹下会生成 engine 文件 yolov8s-cls.FP32.trtmodel 用于模型推理，同时在终端还可以看见模型预测的结果。
OK！以上就是使用 tensorRT_Pro-YOLOv8 推理 YOLOv8-Cls 的大致流程，若有问题，欢迎各位看官批评指正。
结语
博主在这里针对 YOLOv8-Cls 的预处理和后处理做了简单分析，同时与大家分享了 C++ 上的实现流程，目的是帮大家理清思路，更好的完成后续的部署工作😄。感谢各位看到最后，创作不易，读后有收获的看官请帮忙点个👍⭐️
最后大家如果觉得 tensorRT_Pro-YOLOv8 这个 repo 对你有帮助的话，不妨点个 ⭐️ 支持一波，这对博主来说非常重要，感谢各位🙏。
下载链接
参考
-
相关阅读:
【网络安全】-网络安全的分类详解
rauth，一个强大的Python库！
Sedex验厂有证书吗？
利用redis + mysql 完成签到
以太坊合并后展望与机构DeFi的未来
WebGoat 靶场 JWT tokens 四 五 七关通关教程
关于C语言编译环境
A-LOAM学习
穿越时空的创新：解析云原生与Web3.0的奇妙渊源
机械设备经营小程序商城的作用是什么
- 原文地址：https://blog.csdn.net/qq_40672115/article/details/134277392
