YOLOv5多分类识别

最新推荐文章于 2024-08-17 21:40:25 发布

Hello未来

最新推荐文章于 2024-08-17 21:40:25 发布

阅读量268

点赞数 4

文章标签： YOLO 分类 pytorch 深度学习人工智能

本文链接：https://blog.csdn.net/hello000w/article/details/141278982

版权

用于个人记录，好记性不如烂笔头

一、数据集制作

对于多分类来说，是不需要使用labelimg进行打标签的，只需要把分类的图片整理到各个文件夹中即可。

如下图，这里我是做9分类，所以分成9个文件夹。

值得注意的是训练完转成onnx后，文件夹的名字最后会变成索引所对应标签的名字

转成的onnx：

在这里插入图片描述

数据集整理：

在这里插入图片描述

二、训练配置

处理好数据集后，我们还需要去官网下载一个分类的预训练的模型，在这里我使用的是yolov5l-cls.pt作为预训练模型，剩下的就是修改classify文件夹中的train.py即可开始训练.
在这里插入图片描述

配置中设置了训练了5轮，出现下面的图片就可以证明已经启动成功，可以看到只训练了5轮，它的准确率达到了1。训练结束后，会提示Results saved to …\runs\train-cls\exp4，这是存放训练结果的地方

在这里插入图片描述

三、预测配置

训练完后，我们只需要将best.pt模型的路径复制到predict.py对应位置，并修改其他的配置项即可实现预测，我们打开.\runs\train-cls\exp4，将模型路径复制上去

在这里插入图片描述

下面是我对一张图片进行做预测，可以看到图片中左上角有白色的字，那是属于某个类别的概率，图中其他类别都是0，只有1_one的概率为1.00

在这里插入图片描述

四、训练的模型做预测

下面我们将pt模型转成onnx，并使用onnx进行推理。和目标识别一样的，我们只需要修改export.py文件的配置即可。

注意：配置项中–include那一行，一定要自己写上“onnx”，这里我没有截图进去，大概在597行附近

在这里插入图片描述

执行完后会生成onnx，会存放在…runs\train-cls\exp4\weights中
在这里插入图片描述

此时我们打开https://netron.app，将我们的onnx扔进去，我们可以看到输出是shape为(1,9),并且有索引对应识别到的类别，比如0标签就对应1_one这个文件夹。
（1,9）的解释：
可以理解为输入图片，它会输出9个类别的概率，每个类别输出一个概率。

在这里插入图片描述

最后我们随机拿一张图片，将数字2的图片路径写入下方的代码中，它的输出是1，这个1是下标。在上图中，索引1对应的类别是2_two，而2_two里面都是数字2的图片，证明推理成功！

输入的图片：

在这里插入图片描述

onnx_path=r"E:\yolomaster\yolov5\runs\train-cls\exp4\weights\best.onnx"
img_path=r"F:\yolodata\numbers\train\2_two\2.png" 
img=cv2.imread(img_path)
img=img[np.newaxis,:,:,:].astype(np.float32)
img=img.transpose(0,3,1,2)/255.0


provider=ort.get_available_providers()[1 if ort.get_device()=="GPU" else 0]
ort_session=ort.InferenceSession(onnx_path,providers=[provider])

resluts=ort_session.run(output_names=["output"],
                        input_feed={"images":img})

print(resluts[0].shape,type(resluts[0])) # (1, 9) <class 'numpy.ndarray'>

# 找出最大的索引
best_index=resluts[0].argmax(axis=1) 
print(resluts[0])
print(best_index) # [1]