#AI挑战营终点站# rv1106平台使用yolov5识别手写数字

不语arc

#AI挑战营终点站# rv1106平台使用yolov5识别手写数字 [复制链接]

本帖最后由不语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学习。

# 合成MNIST大图
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] # resize的比例尺寸
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)  # 棕色
]

# 加载MNIST数据集
transform = transforms.Compose([
    transforms.ToTensor(),  # 将PIL Image或numpy.ndarray转换为tensor，并缩放到[0.0, 1.0]
    transforms.Normalize((0.5,), (0.5,))  # 归一化到[-1.0, 1.0]
])

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)
# 创建用于保存生成的txt文件的目录
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
    # 用于生成每张图片的label
    curLable = np.zeros((NumOfPic, 5), dtype=float) # 5对应cls x y w h

    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)

    # 将tensor转换回numpy数组
    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  # [C, H, W] to [H, W, C] and scale to 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')

        # 随机resize
        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 # 框的正中心x坐标（归一化）
        curLable[i][2] = (top + img_pil.height / 2) / IMAGE_SIZE # 正中心y坐标
        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:
            # 对于每一行，首先取第一列的整数值，其余列保留浮点数，用空格连接后写入文件
            # 注意：np.int32()用于将浮点数转换为整数，如果是非负数且不需要特别大的整数，np.int_或直接int()也可以
            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指定对应图片的路径。

划分代码如下：

# -*- coding: utf-8 -*-
"""
将所有数据集按照 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为列表，images文件夹下所有后缀为 jpg.jpeg,png 的文件的路径
    total_files = glob((os.path.join(root_path, "images/*[jpg.jpeg,png]")))

    # split
    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)

    # train
    for file in train_files:       # 遍历训练集中所有文件的路径
        ftrain.write(file + "\n")    # 将路径一条一行写入train.txt中，train.txt程序中对应的参数是ftrain
    # val
    for file in val_files:
        fval.write(file + "\n")
    # test
    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

# Classes
nc: 10  # number of classes
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拉流缓冲时间不能设置太小，不然没法看。
模型提升方面。改进方法也只能是改进数据集，人工标注增加样本也不是不行，正所谓人工智能就是有多少人工就有多少智能。