机器视觉抓取（2）——Yolov5_v6.0训练自己的数据集

南栀北辰SDN

已于 2024-04-05 14:11:53 修改

阅读量1.4k

点赞数 24

分类专栏： YOLO系列文章标签： YOLO

于 2023-12-10 21:17:40 首次发布

本文链接：https://blog.youkuaiyun.com/qq_46173016/article/details/134883442

版权

YOLO系列专栏收录该内容

9 篇文章

订阅专栏

前言：在机器视觉抓取（1）中，我们完成了yolov5_v6.0结合D435iF相机实现坐标2d到3d的转变，那么这篇文章将教会你如何训练自己的数据集，识别自己需要的目标，配合机器视觉抓取（1）完成物体定位。

所有文件（包括数据集照片）不能含有中文路径、空格、减号，且数据集照片必须是JPG格式

一. 环境（不能含有中文路径、空格、减号）

本教程使用环境：

pytorch: 1.13.1+cpu

python: 3.7

yolov5_v6.0

ps: 如果要使用GPU，cuda版本要 >=10.1

下载安装yolov5：

yolov5 v6.0官方要求 Python>=3.6 and PyTorch>=1.7

yolov5源码下载：https://github.com/ultralytics/yolov5

下载后，进入pytorch环境进入yolov5文件夹，使用换源的方法安装依赖。

如果你前面安装时没有换源，我强烈建议你使用换源的方法在安装一次

安装过的模块不会在安装，以防缺少模块，影响后续程序运行以及模型训练。

pip install -r requirements.txt -i  https://pypi.tuna.tsinghua.edu.cn/simple

二. 准备工作（文件夹及图片）

1.在 yolov5目录下新建文件夹 VOCData（可以自定义命名）
在这里插入图片描述
2.在VOCData下新建两个文件夹 Annotations 以及 images

images：用于存放要标注的图片（jpg格式）
Annotations ：用于存放标注图片后产生的内容（这里采用XML格式）

三. 使用labelImg标注图片

1. 安装labellmg
labelImg下载：https://github.com/tzutalin/labelImg
不同系统下安装方法都在此链接中有说明，这里只给出 windows +anaconda 进行安装
本人labellmg存放位置
在这里插入图片描述
执行命令前，建议更新一下conda

conda update -n base -c defaults conda

接着执行以下命令：（不要连接外网）

conda install pyqt=5
conda install -c anaconda lxml
pyrcc5 -o libs/resources.py resources.qrc

2. 使用labellmg
在这里插入图片描述
直接运行就可以，也可以在pycharm终端执行

labelImg

勾选 auto save mode 自动保存
![(https://img-blog.csdnimg.cn/direct/e99700ce53fa4a6bbd3c4a0614538e4f.png) 在这里插入图片描述

点击左方边栏或者屏幕右键选择 Create RectBox 即可进行标注。

尽可能的完全拟合标注物体，建议放大标注，也别放大狠了哈。

四. 划分数据集以及配置文件修改

使用pycharm、vscode、python自带的IDLE。如果出现缺少模块的情况（no module named），你可以安装模块，本人使用pycharm

1. 划分训练集、验证集、测试集

在VOCData目录下创建程序 split_train_val.py 并运行
程序如下：（可以不更改）

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')
#数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0  # 训练集和验证集所占比例。 这里没有划分测试集
train_percent = 0.9     # 训练集所占比例，可自己进行调整
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

运行完毕后会生成 ImagesSets\Main 文件夹，且在其下生成测试集、训练集、验证集，存放图片的名字（无后缀.jpg）
在这里插入图片描述

由于没有分配测试集，所以测试集为空。若要分配，更改第 14、15 行代码，更改所在比例即可

2. XML格式转yolo_txt格式

在VOCData目录下创建程序 text_to_yolo.py 并运行
程序如下：
需要将第 7 行改成自己所标注的类别，以及28-30行，59-72行修改各自对于的绝对路径

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val', 'test']
classes = ["apple_A", "apple_B"]  # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def convert_annotation(image_id):
    # 更换成自己的路径
    in_file = open('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/Annotations/%s.xml' % image_id, encoding='UTF-8')
    out_file = open('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/labels/%s.txt' % image_id, 'w')
    out_file = open('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/labels/%s.txt' % image_id, 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        # difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


# 以下路径更改成自己的路径
wd = getcwd()
for image_set in sets:
    if not os.path.exists('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/labels/'):
        os.makedirs('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/labels/')
    image_ids = open('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/ImageSets/Main/%s.txt' % image_set).read().strip().split()

    if not os.path.exists('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/dataSet_path/'):
        os.makedirs('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/dataSet_path/')

    list_file = open('dataSet_path/%s.txt' % image_set, 'w')  # 这行路径不需更改，这是相对路径
    
    for image_id in image_ids:
        list_file.write('F:/Python_Learn_Source/YOLO/yolov5_6.0/VOCData/images/%s.jpg\n' % image_id)
        convert_annotation(image_id)
    list_file.close()

运行后会生成如下 labels 文件夹和 dataSet_path 文件夹。

其中 labels 中为不同图像的标注文件。每个图像对应一个txt文件，文件每一行为一个目标的信息，包括class, x_center, y_center, width, height格式，这种即为 yolo_txt格式
在这里插入图片描述

dataSet_path文件夹包含三个数据集的txt文件，train.txt等txt文件为划分后图像所在位置的路径，如train.txt就含有所有训练集图像的路径。
在这里插入图片描述

3. 配置文件

在 yolov5 目录下的 data 文件夹下新建一个My_Voc.yaml文件（可以自定义命名），用记事本打开。

内容是：

训练集以及验证集（train.txt和val.txt）的路径（可以改为相对路径）

以及 目标的类别数目和类别名称。

注意这里的类别名称顺序，如果训练的结果是反的，这里修改重新训练。

注意这里的类别名称顺序，如果训练的结果是反的，这里修改重新训练。

给出模板： 冒号后面需要加空格

train: D:/Yolov5/yolov5/VOCData/dataSet_path/train.txt
val: D:/Yolov5/yolov5/VOCData/dataSet_path/val.txt

# number of classes
nc: 2

# class names
names: ["apple_A", "apple_B"]

3. 聚类获得先验框
3.1 生成anchors（两种方法）

自动获取anchors

如果目录 yolov5/utils下有 autoanchor.py文件，那么就可以采用自动获取anchors。（yolov5版本偏低是没有的，v6.0是有的）

如果没有此文件。那么就只能手动获取更改anchors

确保 yolov5_6.0/data/hyps/hyp.cratch.yaml中的anchors这行是注释掉的，一般源码没改动的情况下都是注释掉的

采用自动法的话，不用运行，训练时自动调用
如果自动获取anchors的话，训练时只要不添加参数 --noautoanchor（后面会介绍），就会自动运行autoanchor.py。

3.2 修改模型配置文件
选择一个模型，在yolov5目录下的model文件夹下是模型的配置文件，有n、s、m、l、x版本，逐渐增大（随着架构的增大，训练时间也是逐渐增大）。

这里选用 yolov5s.yaml
在这里插入图片描述

复制粘贴备份一份yolov5s原始文件命名为‘yolov5s-原始文件.yaml’

再使用记事本打开 yolov5s.yaml，进行修改

采用自动法获取anchors，只需更改nc 标注类别数，不用更改anchors

五. 模型训练

1. 开始训练

打开yolov5 目录下的 train.py 程序，我们可以多看看这些参数使用。
在这里插入图片描述
常用参数解释如下：
weights：权重文件路径

cfg：存储模型结构的配置文件

data：存储训练、测试数据的文件

epochs：指的就是训练过程中整个数据集将被迭代（训练）了多少次，显卡不行你就调小点。

batch-size：训练完多少张图片才进行权重更新，显卡不行就调小点。

img-size：输入图片宽高，显卡不行就调小点。

device：cuda device, i.e. 0 or 0,1,2,3 or cpu。选择使用GPU还是CPU

workers：线程数。默认是8。

其它参数解释：

noautoanchor：不自动检验更新anchors rect：进行矩形训练

resume：恢复最近保存的模型开始训练

nosave：仅保存最终checkpoint

notest：仅测试最后的epoch

evolve：进化超参数

bucket：gsutil bucket

cache-images：缓存图像以加快训练速度

name：重命名results.txt to results_name.txt

adam：使用adam优化

multi-scale：多尺度训练，img-size +/- 50%

single-cls：单类别的训练集

训练命令如下：

python train.py --weights weights/yolov5s.pt  --cfg models/yolov5s.yaml  --data data/myvoc.yaml --epoch 200 --batch-size 8 --img 640   --device cpu

或者直接在pycharm运行train.py

–weights weights/yolov5s.pt ：这个也许你需要更改路径。我是将yolov5的pt文件都放在weights目录下，你可能没有，需要更改路径。

–epoch 200 ：训练200次

–batch-size 8：训练8张图片后进行权重更新

–device cpu：使用CPU训练。

2. 训练过程

如果你使用GPU训练也有类似下面这个，那是你 cuda 版本不对（不是>=10.1的版本），版本不对无法使用cuda

在这里插入图片描述
训练好的模型会被保存在 yolov5 目录下的 runs/train/weights下

3. 相关问题

如果出现缺少模块的情况（no module named）

回到博客最开始部分，使用换源的方法补充安装yolov5的依赖。

如果出现（页面太小，无法完成操作）的相关问题

那是虚拟内存不足，重新打开页面或者重启电脑试试（这个方法解决的可能性比较低），降低线程 --workes (默认是8) 。最后再试试调小 --batch-size，降低 --epoch

–workers指定为0才成功。

如果都不行，可以看看这个链接 https://product.pconline.com.cn/itbk/software/dnyw/1707/9679137.html

如果训练过程中出现 memory error

那是内存超了，减小 --batch-size 试试，如果还不行降低 --epoch。

将 --epoch 设为100次，–batch-size设为3才成功。
–epoch建议尽量在100次往上吧

重复训练的话，你也许需要将这两个缓存清除掉。

4. 训练可视化

训练时或者训练后可以利用 tensorboard 查看训练可视化
在终端窗口执行：

tensorboard --logdir=runs

在这里插入图片描述

六. 检测效果

使用刚刚训练出的最好的模型 best.pt 来测试，在yolov5_6.0目录下的 runs/train/exp/weights
修改detect.py进行检测
模型在没有标注的数据集上进行推理，在detect.py文件中指定测试图片和测试模型的路径，其他参数（img_size、置信度object confidence threshold、IOU threshold for NMS）可自行修改，如下：
在这里插入图片描述
测试结果保存在 yolov5/runs/detect 目录下

使用下面的命令（该命令中save_txt选项用于生成结果的txt标注文件，不指定则只会生成结果图像），其中，weights使用最满意的训练模型即可，source则提供一个包含所有测试图片的文件夹路径即可。
参考下面链接：

https://blog.youkuaiyun.com/qq_36756866/article/details/109111065?spm=1001.2101.3001.6650.2&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2~default~CTRLIST~default-2.no_search_link

报错问题

在我们运行train.py训练时，不出意外应该会报错：

错误：RuntimeError: result type Float can‘t be cast to the desired output type __int64

这是yolov5_6.0源码与pytorch不同版本兼容性的原因，在这里插入图片描述
解决方法：

将loss.py中gain = torch.ones(7, device=targets.device)改为gain = torch.ones(7, device=targets.device).long()即可。原因是新版本的torch无法自动执行此转换，旧版本torch可以。

在utils文件夹下的loss.py中，ctrl+f搜索gain，找到gain = torch.ones(7, device=targets.device)，将其修改为gain = torch.ones(7, device=targets.device).long()，问题解决。
在这里插入图片描述