YOLOv8多分类识别

文章目录

• • • 1.数据集处理
    • 2.模型训练
    • 3.模型推理
    • 4.一些问题
    • • （1）V8它是怎么知道我们训练了多少类
    • 5.参数说明

1.数据集处理

​ v8的数据处理方式和v5是一样的，不需要打标签，只需要将数据整理好后放在各个文件夹中，文件夹的名称就是它的标签名，这个可以通过onnx可以看到。

训练后转成onnx如下

这个onnx是我之前训练好的模型并转成了onnx

数据整理如下

每个文件夹里面存放的都是320x320的图片，这里我做了9个分类，每个类别存放在不同的文件夹中

2.模型训练

​ 整理好数据集后，就可以开启训练了，跟V5一样，只是V8可以不用下载工程文件，安装ultralytics这个包后调用，并调整参数即可，参数说明放在第4点的参数说明那里

注意：记得下载预训练的模型，不然运行代码的时候它再自动下载的话会比较慢

from ultralytics import YOLO

if __name__ == '__main__':
    model_path = r"E:\yolodata\yolov8s-cls.pt"
    model=YOLO(model_path)
    model.train(
        # 数据文件路径，例如 coco128.yaml，分类任务直接写数据集文件夹路径即可
        data=r"E:\yolodata\numbers",
        imgsz=320, # 图像大小
        device=0, # 运行设备
        batch=32, # 每批次的图像数量 （-1 表示自动批次）
        epochs=10) # 训练的周期数

​ 当执行成功以上代码后，会在该代码文件的同一文件夹下生成一个runs文件夹，最里面的weight文件夹会存放我们训练好的pt模型。训练完后会提示我们所保存的位置，到时候记得看打印台即可.

​

​ 我们打开训练代码的文件夹，会发现多了一个runs文件，预训练的模型是我自己提前下载好的，numbers是我们的数据集，v_8train.py是我训练代码的文件

​

​ 打开runs文件夹里面的train文件，里面存放我们训练的结果已经存放模型的weight文件夹

​ 我们可以随便打开一个验证的图片看看，虽然只训练了10轮，可以看到效果还行，没有很差。

3.模型推理

​ 有了pt模型后，我们就可以转成onnx方便后续部署

from ultralytics import YOLO

if __name__ == '__main__':
    model_path = r"E:\yolodata\runs\classify\train\weights\best.pt"
    model=YOLO(model_path)
    model.export(format="onnx",imgsz=320)

​ 从onnx我们知道它的输入是tensor: float32[1,3,320,320]，输出是tensor: float32[1,9]，剩下我们只需要导入onnx模型，然后将图片进行处理成它输入要的shape，并进行归一化处理即可

​ 通过上面的onnx图片，我们知道names部分包含标签的名字以及它的索引（下标），在代码中我的图片是一张3的图片，它对应的标签名是**“3_three”，所以可以通过上面的onnx知道它的索引是2**，所以只要代码中输出是2则表示推理成功。

import cv2
import os
import onnxruntime as ort
import numpy as np



onnx_path=r"E:\yolodata\runs\classify\train\weights\best.onnx"
provider = ort.get_available_providers()[1 if ort.get_device()=="GPU" else 0]

ort_session=ort.InferenceSession(onnx_path, providers=[provider])

img_path=r"E:\yolodata\numbers\val\3_three\26.png"
img_src=cv2.imread(img_path)
img_dst=img_src[np.newaxis, :,:,:]/255.
img_dst=np.transpose(img_dst, (0,3,1,2)).astype(np.float32)
# print(img_dst.shape,img_dst.dtype) # (1, 3, 320, 320) float32


output=ort_session.run(output_names=['output0'], input_feed={'images':img_dst})

best_index=np.argmax(output[0],axis=1)

print(best_index) # 输出[2] 推理成功

4.一些问题

（1）V8它是怎么知道我们训练了多少类

那是因为它的数据加载中用了torchvision.datasets.ImageFolder，

会自动根据子目录名作为类别标签，并将每个子目录下的图像作为该类别的样本。
# 同时，它还会返回一个映射关系，表示类别标签到实际类别的字典，可以通过class_to_idx属性访问

在代码中的这个位置

5.参数说明

model.train()参数说明

参考链接：https://blog.csdn.net/qq_42452134/article/details/135168930

model.export()参数说明

参考链接：https://blog.csdn.net/qq_42761751/article/details/140825169

完结！！！