本帖最后由 不语arc 于 2024-6-5 15:36 编辑
背景:
使用前两期训练的 针对MNIST数据集的 模型存在以下缺点:
- 图片分辨率过低28x28,与现实存在较大差异。
- 对摄像头采集到的大图一次只能识别一个数字。
- 前处理操作 如找到大图中轮廓最大的数字、对送入模型的小图进行二值化 等操作注定了识别场景有限。
为了解决以上问题(并且luckfox已经有使用yolov5的demo工程了)所以决定使用yolov5完成手写数字的识别。
效果预览:
数据集准备
yolov5的输入图像分辨率是640x640,而MNIST数据集的尺寸是28x28,应该如何解决?
如果直接将MNIST小图拼接在一起,就好像这样,特征过于明显,颜色都是黑底白字、目标框的位置都是固定的。数据集容易过拟合。
想法:使用MNIST 28x28分辨率的小图 合成640x640的大图。大图包含若干小图,同时输出对应label
效果预览:
目标:数字的颜色随机(增强数据),大图背景随机并且小图背景(数字外的区域)跟随、小图的尺寸随机被resize、小图的位置随机。自动生成每张小图的label标签。
过程:
- 对图片,加载MNIST数据集(定义每张大图有12个小图)、创建一个640x640的背景(颜色随机)、对每一张小图有:检测数字位置与背景位置、改变对应位置的像素值、随机resize小图、随机摆放、最后将小图粘贴到背景上。
- 对标签label,可以大概设定每一张小图的预测框是28x28,为了避免预测框都是正方形,代码里对长宽都randint(20,28),因为20x20的框也差不多能框住这个小图了。再根据上面的resize值、随机拜访位置值求出对应的label。yolo系列的标签一个框对用5个值,分别是类别、中心点坐标x、中心点坐标y、宽w、高h。后四个指标是归一化之后的float。具体知识可以搜索yolo学习。
-
- import torch
- import torchvision
- import torchvision.transforms as transforms
- from PIL import Image
- import numpy as np
- import random
- import os
-
-
- NUM_IMAGES_TO_GENERATE = 400
- IMAGE_SIZE = 640
- Resize_ratios = [1, 1.5, 2, 2.5, 3, 3.5, 4]
- NumOfPic = 12
- backcolors = [
- (255, 0, 0),
- (0, 255, 0),
- (0, 0, 255),
- (255, 255, 0),
- (0, 255, 255),
- (255, 0, 255),
- (128, 128, 128),
- (255, 255, 255),
- (0, 0, 0),
- (210, 105, 30)
- ]
-
-
- transform = transforms.Compose([
- transforms.ToTensor(),
- transforms.Normalize((0.5,), (0.5,))
- ])
-
- trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
- trainloader = torch.utils.data.DataLoader(trainset, batch_size=NumOfPic, shuffle=True)
-
-
- images_dir = 'dataset\\images'
- if not os.path.exists(images_dir):
- os.makedirs(images_dir)
-
- labels_dir = 'dataset\\labels'
- if not os.path.exists(labels_dir):
- os.makedirs(labels_dir)
-
-
- image_counter = 0
-
-
- for data in trainloader:
- if image_counter >= NUM_IMAGES_TO_GENERATE:
- break
-
- curLable = np.zeros((NumOfPic, 5), dtype=float)
-
- img_tensor, label = data
-
- label_array = label.numpy()
- curLable[:, 0] = label_array
-
- backcolorNum = random.randint(0, 9)
- backcolor = backcolors[backcolorNum]
-
- background = Image.new('RGB', (IMAGE_SIZE, IMAGE_SIZE), color=backcolor)
-
-
- rgb_tensor = img_tensor.repeat(1, 3, 1, 1)
- for i, img in enumerate(rgb_tensor):
- img_np = img.numpy().squeeze().transpose(1, 2, 0) * 255
-
-
- colored_image = np.zeros((img_np.shape[0], img_np.shape[1], 3), dtype=np.uint8)
-
-
- while True:
- numrandom = random.randint(0, 9)
- if numrandom != backcolorNum:
- break
- colors = backcolors[numrandom]
-
- white_pixel_indices = np.where(img_np > 0)
- black_pixel_indices = np.where(img_np == -255)
-
- colored_image[white_pixel_indices[0], white_pixel_indices[1]] = colors
- colored_image[black_pixel_indices[0], black_pixel_indices[1]] = backcolor
- img_pil = Image.fromarray(colored_image, 'RGB')
-
-
- random_ratio = random.randint(0, len(Resize_ratios)-1)
- img_pil = img_pil.resize((int(28*Resize_ratios[random_ratio]), int(28*Resize_ratios[random_ratio])))
- w_random = random.randint(20, 28)
- h_random = random.randint(21, 28)
- curLable[i][3] = w_random/IMAGE_SIZE*Resize_ratios[random_ratio]
- curLable[i][4] = h_random/IMAGE_SIZE*Resize_ratios[random_ratio]
-
-
- left = random.randint(0, IMAGE_SIZE - img_pil.width)
- top = random.randint(0, IMAGE_SIZE - img_pil.height)
- curLable[i][1] = (left+img_pil.width/2)/IMAGE_SIZE
- curLable[i][2] = (top + img_pil.height / 2) / IMAGE_SIZE
- background.paste(img_pil, (left, top))
-
-
- output_img_path = os.path.join(images_dir, f'image_{image_counter}.jpg')
- background.save(output_img_path)
-
-
- txtPath = os.path.join(labels_dir, f'image_{image_counter}.txt')
- with open(txtPath, 'w') as file:
-
- for row in curLable:
-
-
- line = ' '.join([str(int(row[0]))] + [f'{x:.6f}' for x in row[1:]])
- file.write(line + '\n')
-
-
-
- image_counter += 1
-
-
- print(f'Generated {image_counter} images and labels.')
生成了200张图片及标签(其实还可以生成很多张),除此之外,我加了一百二十多张通用图片,自行标注了几十张图片,具体看附件。将他们按照8:1:1划分训练集、验证集、测试集。生成对应的txt指定对应图片的路径。
划分代码如下:
-
- """
- 将所有数据集按照 8:1:1的比例划分为训练集、验证集、测试集。输出为txt文本
- 使用方法:指定数据集的主目录为root_path
- """
- from glob import glob
- import os
- from sklearn.model_selection import train_test_split
- import random
- import shutil
- root_path = r'D:\Users\Yzhe\PycharmProjects\MNIST_Test\dataset'
-
-
- def toTrainValTest(root_path):
- ftrain = open(root_path + '/train.txt', 'w')
- fval = open(root_path + '/val.txt', 'w')
- ftest = open(root_path + '/test.txt', 'w')
-
-
- total_files = glob((os.path.join(root_path, "images/*[jpg.jpeg,png]")))
-
-
- train_files, val_test_files = train_test_split(total_files, test_size=0.2, random_state=42)
-
- val_files, test_files = train_test_split(val_test_files, test_size=0.5, random_state=42)
-
-
- for file in train_files:
- ftrain.write(file + "\n")
-
- for file in val_files:
- fval.write(file + "\n")
-
- for file in test_files:
- ftest.write(file +"\n")
-
- ftrain.close()
- fval.close()
- ftest.close()
-
- toTrainValTest(root_path)
-
生成了对应的3个txt。txt内容是这样的:
至此数据集准备工作完成。
yolov5训练
yolov5训练。(这部分仍然在windows完成,需要显卡。我是4060,8G显存)
环境配置部分网上教程挺多的,使用的anaconda+pycharm。使用anaconda创建了一个torch2的环境,这个环境满足运行yolov5.
- 下载yolov5代码(7.0版本)
https://github.com/ultralytics/yolov5/tree/v7.0
- 配置
2.1 在data目录下新建一个MNIST目录,在该目录下放置刚才生成的txt。
2.2 在此目录下新建一个yaml文件,内容如下:
- train: data/MNIST/train.txt
- val: data/MNIST/val.txt
- test: data/MNIST/test.txt
-
-
- nc: 10
- names: ['0','1','2','3','4','5','6','7','8','9']
2.3 在train.py文件下,找到parse_opt函数(用于指定参数)
照着这个配置即可,如果显存不大就调小batchsize。epochs是训练1000轮,实际上超过一百轮没提升就提前结束了。
训练结果保存在runs/train/exp目录下,权重文件选择weights/best.pt。训练过程如下。
- 导出为onnx模型(这部分参考https://wiki.luckfox.com/zh/Luckfox-Pico/Luckfox-Pico-RKNN-Test/#4-yolo-v5)还是在windows操作。
3.1 获取yolov5代码(这里用的是rockchip的代码,他在原yolov5的基础上增加了对rknn的支持)
git clone https://github.com/airockchip/yolov5.git
具体改动是:
3.2 激活第二步使用的环境torch2:conda activate torch2
3.3 将刚才训练好的best.pt复制到项目根目录下,并重命名为yolov5s.pt
3.4 在yolov5/python目录下,运行export.py
python export.py --rknpu --weight yolov5s.pt.
在根目录下生成了yolov5s.onnx
3.5 使用netron网站可以观察模型结构。
https://netron.app/
网站查看onnx模型结构
- 转rknn。这一步需要在Ubuntu环境下,使用前两期训练营配置的toolkit2_0.0环境
4.1 下载 rknn_model_zoo
git clone https://github.com/airockchip/rknn_model_zoo.git
4.2 将刚刚得到的ONNX模型文件拷贝至 rknn_model_zoo/examples/yolov5/model 目录
4.3 修改数据集内容(因为是自行训练的数据集,不是coco数据集),这一步是为了量化。
4.3.1 修改rknn_model_zoo/examples/yolov5/python/convert.py文件
4.3.2 在rknn_model_zoo/datasets/COCO目录下新建一个文件夹、一个txt。
4.3.3 txt内容如下:(用于指定数据集位置)
- ./myMNIST/image_1.jpg
- ./myMNIST/image_2.jpg
- ./myMNIST/image_3.jpg
- ./myMNIST/image_4.jpg
- ./myMNIST/image_5.jpg
- ./myMNIST/image_6.jpg
- ./myMNIST/image_7.jpg
- ./myMNIST/image_8.jpg
- ./myMNIST/image_9.jpg
- ./myMNIST/image_10.jpg
- ./myMNIST/image_11.jpg
- ./myMNIST/image_12.jpg
- ./myMNIST/image_13.jpg
- ./myMNIST/image_14.jpg
- ./myMNIST/image_15.jpg
4.4.4 在myMNIST目录下复制指定15张图片。(图片数量、具体哪一张图片应该是随意的)
4.4 执行 rknn_model_zoo/examples/yolov5/python 目录下的模型转换程序 convert.py,使用方法:
- conda activate toolkit2_0.0
- cd <实际为准>/rknn_model_zoo/examples/yolov5/python
- python3 convert.py ../model/yolov5s.onnx rv1106
在rknn_model_zoo/examples/yolov5/model目录下生成了yolov5.rknn
- 配置最终运行的项目
5.1 下载luckfox的程序
git clone https://github.com/luckfox-eng29/luckfox_pico_rtsp_yolov5.git
5.2 修改model目录内容
5.2.1 替换刚才生成的yolov5.rknn文件
5.2.2 修改coco_80_labels_list.txt文件
其内容是:用于指定十种类别。
5.3 修改 include/rknn/postprocess.h的内容(这个bug找了一晚上),将类别数量改成10
5.4 编译:
- export LUCKFOX_SDK_PATH=<Your Luckfox-pico SDK Path>
- mkdir build
- cd build
- cmake ..
- make && make install
5.5 将项目目录下的luckfox_rtsp_yolov5_demo上传到开发板。
5.6 开发板运行
在luckfox_rtsp_yolov5_demo目录下,
chmod 777 luckfox_rtsp_yolov5
./luckfox_rtsp_yolov5
5.7 使用vlc拉流查看摄像头画面。rtsp://172.32.0.93/live/0。这里我直接拍摄了LCD屏幕效果(LCD显示看我上一次分享帖子)
分析:
在检测精度上,检测结果已经相当可以了,不会乱框,性能的提升与数据集强相关。
存在的问题:
- 速度确实慢了,明显感觉比之前慢很多,而且vlc拉流缓冲时间不能设置太小,不然没法看。
- 模型提升方面。改进方法也只能是改进数据集,人工标注增加样本也不是不行,正所谓人工智能就是有多少人工就有多少智能。