当前位置：物联沃-IOTWORD物联网 > 技术教程 > 使用YOLOV5训练自己的数据集（以王者荣耀为例）
代码收藏家代码收藏家 技术教程

使用YOLOV5训练自己的数据集（以王者荣耀为例）

注：本文只是记录笔者使用yolov5训练自己的数据集的实现过程，不讲原理，如果想学原理请移步点击量比较高的大佬博客，笔者只是大二本科生，如有问题还请多多指教!

一，什么是yolov5？

“YOLO”是一个对象检测算法的名字，YOLO将对象检测重新定义为一个回归问题。它将单个卷积神经网络(CNN)应用于整个图像，将图像分成网格，并预测每个网格的类概率和边界框。YOLO非常快。由于检测问题是一个回归问题，所以不需要复杂的管道。它比“R-CNN”快1000倍，比“Fast R-CNN”快100倍。YOLOv5是YOLO的第五个版本，目前最新的Tags已经到第六个版本了。具体技术原理可到官网了解。
官网地址：https://github.com/ultralytics/yolov5

二，准备工作

1.环境的配置
具体步骤可参考其他其他博主博客（关键词：YOLOv5环境配置）
以下是官方给出的配置要求：
1）PyTorch版本，官方要求torch >= 1.7.0,笔者用的是1.12.0
2）Python版本，官方要求python >=3.8，笔者是3.9.6 ，命令行输入python --version即可查看自己的python版本

3）cuda版本，大家根据自己的设备硬件条件选择合适的即可，笔者用的是10.2
2.yolov5的下载：
进入GitHub官方网站后点击Download ZIP就可以下载到本地了

下载到本地后，将压缩包进行解压，解压后可以根据requirements.txt直接查看其所需的依赖包，如果你缺少某个包可以直接通过pip 来安装。
——–在你配置完YOLOV5的环境后并且YOLOV5压缩包已经成功下载到本地的条件下，可以直接运行文件夹里面的detect.py，他是来帮助你测试你的环境配置是否OK，运行成功后，会在终端显示结果存放在那个位置，不出意外的话都是在runs/detect/exp中能发现被处理过的图片，如下图所示即代表运行成功，若程序报错，说明你的环境配置存在问题，可能是某些依赖包的版本不符，可根据问题的描述，自行百度解决。

三，数据集的准备

首先是你的手头上要有至少上百张的图片以用来训练，图片越多，训练效果越好。笔者以王者荣耀为素材，来实现对目标的检测。资源和效果展示放在文末。
第一步，在yolov5的data文件夹下创建以下文件夹：

以上文件夹中，images存放你的原始图片，Annotations存放标注的xml文件（接下来会讲解怎么生成xml文件），JPEGImages存放的和images文件的内容一样，也是原图片（直接将images文件夹的内容复制过来即可）
第二步，完成文件夹的创建以后，我们就要使用工具来生成xml文件，这里我们使用的是labelme这个工具，
官网地址：https://github.com/wkentaro/labelme，我们进入到官网地址后将压缩包下载到本地进行解压
解压完成后无需配置任何环境，进入文件夹中按住shift键的同时鼠标右键，点击在此处打开power shell窗口（与命令行窗口一样，可以运行命令），在打开的窗口中输入 python labelimg.py 回车

将会出现一下界面，这就是我们打标签（生成xml文件）所使用的界面

界面中，点击左侧的打开目录，选择第一步你创建的文件夹中的images文件夹（你原始图片存放的文件夹）
然后点击改变存放目录，选择你创建的Annotations文件夹。最后点击上方的查看，选择自动保存模式。
接下来就是对你的原始图片进行标注，快捷键：w:创建标注框,d：下一张,a:上一张, 创建标注框后，命名后点击OK然后下一张即可
对你所有的原始图片打完标签后，你的所有xml文件都已经生成，存放在Annotations文件夹中，接下来，在yolov5的根目录下创建一个trans_txt.py文件,文件代码如下：

import os
import random
trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'data/Annotations'
txtsavepath = 'data/ImageSets'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
ftrainval = open('data/ImageSets/trainval.txt', 'w')
ftest = open('data/ImageSets/test.txt', 'w')
ftrain = open('data/ImageSets/train.txt', 'w')
fval = open('data/ImageSets/val.txt', 'w')
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

继续创建voc_label.py文件，代码如下：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = ['train', 'test','val']
classes = ['heros','soliders','tower','red_buff','blue_buff','pig','wolf','bird','lizards']
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    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('data/Annotations/%s.xml' % (image_id))
    out_file = open('data/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
        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))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
print(wd)
for image_set in sets:
    if not os.path.exists('data/labels/'):
        os.makedirs('data/labels/')
    image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()
    list_file = open('data/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write('data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

注意！！！这个文件需要修改，代码中
classes = ['heros','soliders','tower','red_buff','blue_buff','pig','wolf','bird','lizards']这一条语句需要根据你打标签的时候的类别进行修改，有几类，就写几类！！！！
完成修改后，依次运行两个文件，成功运行后：
（1） labels下生成txt文件（显示数据集的具体标注数据）
（2） ImageSets下生成四个txt文件

（3） data下生成三个txt文件（带有图片的路径）

四，yaml文件的修改

在data文件夹中找到coco.yaml文件，复制一份，任意命名，我这里命名为wangzhe.yaml，还是存放在data文件夹下即可，然后对wangzhe.yaml的部分内容进行修改，我这里是使用Pycharm打开的yaml文件，也可以使用记事本打开，需要对以下内容需要修改

！！！！！！！！注意斜杠的方向，不要弄错了！！！！！！！

  1. 修改train,val,test的路径为自己刚刚生成的路径（data目录下面的这三个train,val,test 文件的路径）
  2. nc 的数字代表数据集的类别，我这里有9类，所以修改为9
  3. names 里为自己数据集打标签时候所有类的名称，我这里是：names: ['heros','soliders','tower','red_buff','blue_buff','pig','wolf','bird','lizards']

五，models模型文件的修改

yolov5有4种配置，不同配置的特性如下，我这里使用的yolov5s，因为其速度最快，但是效果不是很好

找到models文件夹下面的yolov5s.yaml文件打开只需修改一个地方如下：只需要将nc 修改为你的类别数目即可

六，进行训练

首先，在yolov5文件夹下创建一个名为weights的文件夹，然后把data/scripts文件下download_weights.sh放在weights文件下，然后将权重文件yolov5s.pt放到weights文件夹中，在官网地址下滑找到对应权重文件下载即可，但是官方下载非常慢，所以我在这里分享给大家
链接：https://pan.baidu.com/s/1fRIntAeqVf3yNE1GpTBVLw?pwd=s13z
提取码：s13z
下载后将权重文件放到weights中即可，如下：

训练前需要对train.py进行内容的修改，打开train.py，直接找到最下面一堆参数的地方，对以下内容进行修改（修改default）
1.–weights:即刚才下载的权重文件路径，如果你是根据我刚才所讲的在yolov5文件夹下面创建的weights文件夹的话，就修改为上图与我一样的路径就OK
2.–cfg: 我们上面说到我们使用yolov5s模型，所以路径就是models文件夹下面的yolov5s.yaml，如果 你使用的是其他模型，修改为models文件下的对应模型的yaml文件即可
3.–data: 这个路径就选我们在第四个步骤中自己配置的yaml文件，我的是wangzhe.yaml
4.–epochs:训练过程中整个数据集将被迭代多少次，显卡不行的将这个数调小一点，根据自己设备硬件情况来设定
5.–batch-size：一次看完多少张图片才进行权重更新，显卡不行需要调小点。最好是2的倍数，（如2，4，6，8…）
———————————————————————————————————————————–
如果epochs和batch-size设置不好的话，在训练的时候会出现cuda内存溢出的问题，即内存占用过高电脑带动不了，跑不了，如果训练的时候报错，错误描述是关于内存的，就说明这两个参数设置的不合适，适当调小一点即可

这些参数设置完以后运行train.py即开始训练，若出现问题，可自行百度解决，如果以上参数和路径设置正确的话，基本上不会出现问题的，下图就是训练过程：
运行结束后会告诉你把结果保存在哪个文件下，根据运行结果找到文件即可。
一般而言，训练好的模型会被保存在yolov5目录下的D:\yolov5-master\runs\train|\exp下面的weights文件夹中，会有best.pt和last.pt两个模型。best.pt是指在你训练过程中得到的效果最佳的模型，last.pt是指在你训练结束后最终得到的模型，在使用模型进行验证时，我们往往选用best.pt

七，测试detect.py

同样的我们需要打开detect.py进行某些参数的修改，打开detect.py，直接找到最下面一堆参数的地方，对以下内容进行修改
–weights:这里就要使用我们刚才训练得到的best.pt，路径修改为best.pt的存放路径即可
–source: 这里就是选择我们需要测试的图片或视频，你的图片或者是视频放在那里，就写对应的路径，如果想调用我们的电摄像头，直接设置为0即可

修改完参数后直接运行detect.py就可以了，程序运行结束会告诉你测试结果放在那里，找到对应的文件夹即可看到。

八，总结

到这里，整个过程就算结束了，在实现的过程中笔者遇到了很多问题，这些问题都可以百度解决，我相信你们也可以，当然我希望你们不会出错 哈哈，有问题欢迎在评论区留言，如若文章哪里有错误，还请指出，一定虚心学习！！
下面我将我的数据集分享在这里，以及我的结果展示：
原视频：链接：https://www.bilibili.com/video/BV1AW4y1S7MK/?vd_source=6cbf6f3c5a42f6acc6f320976ad01fb4

yolov5检测原视频

检测视频效果展示:链接：https://www.bilibili.com/video/BV1tZ4y1Y7j4/?vd_source=6cbf6f3c5a42f6acc6f320976ad01fb4&t=2.3

检测视频

数据集分享(原图片和原始视频):
链接：https://pan.baidu.com/s/1sFQ3PZES8OCyjcNrOCcGYA?pwd=vgy2
提取码：vgy2

您可能感兴趣的内容:

物联沃分享整理
物联沃-IOTWORD物联网 » 使用YOLOV5训练自己的数据集（以王者荣耀为例）

代码收藏家 普通

分享到：

相关推荐

发表评论