用于个人记录,好记性不如烂笔头
一、数据集制作
对于多分类来说,是不需要使用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]
完结!!!