1、简介
本文主要详细对YOLOV5网络训练自己的数据集进行讲解，从基础网络下载，到训练自己的数据集，详细的对操作过程进行讲解，适合小白学习以及老手回顾。
2、YOLOV5网络下载
2.1 网络查找
首先点击项目地址：YOLOV5
2.2 版本选择
1.点击tags进行版本选择，如下图所示：
2.选择所需的版本，本次以V6.1版本进行介绍，如下图所示：
3.返回主页进行下载，如下图所示：
在主界面可以看到以及切换到V6.1版本，接着进行下载，如下图所示：
3 环境搭建
将下载好的网络打开，找到requirements.txt文件点击打开，如下图所示：
接着按照requirements.txt文件要求在终端输入pip install -r requirements.txt，进行各个包搭建，如下图所示：
接着选择detect.py运行，如下所示：
从报错可知需要进行yolov5s.pt进行下载，下载链接如下：
链接：https://pan.baidu.com/s/1dbALmj7AWUQsVz6GZL9Kiw
提取码：gp2w
将下载好的yolov5s.pt文件拽入总文件夹下，在此运行detect.py，可以进行检测，如下所示：
4 训练集准备
在主目录下新建文件夹VOCData，如下所示：
接着在VOCData文件夹下新建Annotations、images文件夹，如下所示：
images文件夹中存放JPG图片,如下所示：
在images文件夹中建立classes.txt 定义自己要标注的所有类别，如下图所示：
Annotations ：存放标注的标签文件。
5 labelImg标注图片
图片标注点击：利用labelimg制作自己的数据集。
6划分训练集、验证集、测试集
6.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 文件夹，并且Main文件夹下会生成四个子文件如下图所示：
6.2在VOCData目录下创建程序 xml_to_yolo.py 并运行
新建文件如下图所示：
文件代码如下所示：
# -*- coding: utf-8 -*- import xml.etree.ElementTree as ET import os from os import getcwd sets = ['train', 'val', 'test'] classes = ["windows"] # 改成自己的类别 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('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/VOCData/%s.xml' % (image_id), encoding='UTF-8') out_file = open('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/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('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/VOCData/labels/'): os.makedirs('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/VOCData/labels/') image_ids = open('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/VOCData/ImageSets/Main/%s.txt' % (image_set)).read().strip().split() if not os.path.exists('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/VOCData/dataSet_path/'): os.makedirs('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/VOCData/dataSet_path/') list_file = open('dataSet_path/%s.txt' % (image_set), 'w') # 这行路径不需更改，这是相对路径 for image_id in image_ids: list_file.write('D:/Test/Desktop/Desktop/TSET/v5mobilenetsmall/yolov5-master/VOCData/images/%s.jpg\n' % (image_id)) convert_annotation(image_id) list_file.close()
第七行类别需要修改自己的类别，如下图所示：
第27、28、57、58、59、61、62、67行改成自己的绝对路径如下图所示：
如果本身有的数据集就是yolo格式，那么可以使用如下代码进行yolo转xml:
from xml.dom.minidom import Document import os import cv2 # def makexml(txtPath, xmlPath, picPath): # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径 def makexml(images, Annotations, xmlPath): # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径 """此函数用于将yolo格式txt标注文件转换为voc格式xml标注文件 在自己的标注图片文件夹下建三个子文件夹，分别命名为picture、txt、xml """ dic = {'0': "boat", # 创建字典用来对类型进行转换 '1': "cat", # 此处的字典要与自己的classes.txt文件中的类对应，且顺序要一致 } files = os.listdir(txtPath) for i, name in enumerate(files): xmlBuilder = Document() annotation = xmlBuilder.createElement("annotation") # 创建annotation标签 xmlBuilder.appendChild(annotation) txtFile = open(txtPath + name) txtList = txtFile.readlines() img = cv2.imread(picPath + name[0:-4] + ".jpg") Pheight, Pwidth, Pdepth = img.shape folder = xmlBuilder.createElement("folder") # folder标签 foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset") folder.appendChild(foldercontent) annotation.appendChild(folder) # folder标签结束 filename = xmlBuilder.createElement("filename") # filename标签 filenamecontent = xmlBuilder.createTextNode(name[0:-4] + ".jpg") filename.appendChild(filenamecontent) annotation.appendChild(filename) # filename标签结束 size = xmlBuilder.createElement("size") # size标签 width = xmlBuilder.createElement("width") # size子标签width widthcontent = xmlBuilder.createTextNode(str(Pwidth)) width.appendChild(widthcontent) size.appendChild(width) # size子标签width结束 height = xmlBuilder.createElement("height") # size子标签height heightcontent = xmlBuilder.createTextNode(str(Pheight)) height.appendChild(heightcontent) size.appendChild(height) # size子标签height结束 depth = xmlBuilder.createElement("depth") # size子标签depth depthcontent = xmlBuilder.createTextNode(str(Pdepth)) depth.appendChild(depthcontent) size.appendChild(depth) # size子标签depth结束 annotation.appendChild(size) # size标签结束 for j in txtList: oneline = j.strip().split(" ") object = xmlBuilder.createElement("object") # object 标签 picname = xmlBuilder.createElement("name") # name标签 namecontent = xmlBuilder.createTextNode(dic[oneline[0]]) picname.appendChild(namecontent) object.appendChild(picname) # name标签结束 pose = xmlBuilder.createElement("pose") # pose标签 posecontent = xmlBuilder.createTextNode("Unspecified") pose.appendChild(posecontent) object.appendChild(pose) # pose标签结束 truncated = xmlBuilder.createElement("truncated") # truncated标签 truncatedContent = xmlBuilder.createTextNode("0") truncated.appendChild(truncatedContent) object.appendChild(truncated) # truncated标签结束 difficult = xmlBuilder.createElement("difficult") # difficult标签 difficultcontent = xmlBuilder.createTextNode("0") difficult.appendChild(difficultcontent) object.appendChild(difficult) # difficult标签结束 bndbox = xmlBuilder.createElement("bndbox") # bndbox标签 xmin = xmlBuilder.createElement("xmin") # xmin标签 mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth) xminContent = xmlBuilder.createTextNode(str(mathData)) xmin.appendChild(xminContent) bndbox.appendChild(xmin) # xmin标签结束 ymin = xmlBuilder.createElement("ymin") # ymin标签 mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight) yminContent = xmlBuilder.createTextNode(str(mathData)) ymin.appendChild(yminContent) bndbox.appendChild(ymin) # ymin标签结束 xmax = xmlBuilder.createElement("xmax") # xmax标签 mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth) xmaxContent = xmlBuilder.createTextNode(str(mathData)) xmax.appendChild(xmaxContent) bndbox.appendChild(xmax) # xmax标签结束 ymax = xmlBuilder.createElement("ymax") # ymax标签 mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight) ymaxContent = xmlBuilder.createTextNode(str(mathData)) ymax.appendChild(ymaxContent) bndbox.appendChild(ymax) # ymax标签结束 object.appendChild(bndbox) # bndbox标签结束 annotation.appendChild(object) # object标签结束 f = open(xmlPath + name[0:-4] + ".xml", 'w') xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8') f.close() if __name__ == "__main__": picPath = "D:/Desktop/TSET/v5-6.1/VOCData/images/" # 图片所在文件夹路径，后面的/一定要带上 txtPath = "D:/Desktop/TSET/v5-6.1/VOCData/Annotations/" # txt所在文件夹路径，后面的/一定要带上 xmlPath = "D:/Desktop/TSET/v5-6.1/VOCData/xmlPath/" # xml文件保存路径，后面的/一定要带上 makexml(images, Annotations, xmlPath)
仅仅需要在下方三个路径进行修改即可：
接着将生成的xml文件复制到Annotations文件夹下：
7 文件配置
7.1在data目录下新建一个XX.yaml文件
分别对训练集、验证集路径进行设置以及类别个数、名称进行修改，如下所示：
代码如下所示：
train: D:/Desktop/TSET/v5-6.1/VOCData/dataSet_path/train.txt #训练集 val: D:/Desktop/TSET/v5-6.1/VOCData/dataSet_path/val.txt #验证集 # number of classes nc: 1 # 类别个数 # class names names: ["windows"] # 类别
7.2聚类获得先验框
检查一下是否有下图文件(低版本是没有的，没有的需要自行手动获取)
接着选中下图的yolov5s.yaml文件对类别个数进行修改，如下图所示：
接着在models文件下将yolov5s.yaml文件复制重命名windows，如下所示：
8开始训练
点击train.py文件仅仅需要对下图参数进行修改：
接着对训练轮数进行设置，如下所示：
运行后，会发现下图错误：
说明虚拟内存不足，utils路径下找到datasets.py这个文件，将119行num_workers=0,如下所示：
接着遇到RuntimeError: result type Float can't be cast to the desired output type __int64问题，如下所示：
loss.py在utils文件夹下，ctrl+f搜索gain，找到gain = torch.ones(7, device=targets.device)，将其修改为gain = torch.ones(7, device=targets.device).long()，问题解决
接着点击运行，开始训练，如下图所示：
9 启用tensorbord查看参数
首先打开pycharm的命令控制终端，输入如下命令：
tensorboard --logdir=runs/train
接着将得到得网址复制，如下图所示：
点击网址，就会跳转到浏览器中，如下图所示：
上文如有错误，恳请给位大佬指正。
