目录
将划分完的测试训练集xml文件转为yolov5可识别的数据集
导言
发这篇文章的目的不仅仅是教会小白也可以使用模型训练，拥有自己的模型，也算是对我这段时间的一个学习的记录吧，万一之后某一天还需要用到的时候可以再次捡起来进行学习和操作。本文采用Python 3.8 版本，使用的是Pycharm 进行教学，文中用到的库，在文件配置有提及。
更新时间：2022 年 01月 26日 21：08
更新时间：2022 年 02月 02日 12：00
更新时间：2022 年 11月 25日 11：14
更新了被和谐的模型下载链接
前期的准备
下载机器学习训练源码
1、YOLOv5的代码链接（大概4M）：
GitHub - ultralytics/yolov5: YOLOv5 🚀 官方下载（推荐）：
https://github.com/ultralytics/yolov5
环境：
Python版本：官方文档给出的是Python3.8，但是其他博主说3.7左右也可以具体没试过
运行yolov5的包都在代码内的 requirements.txt 文件下，缺什么装什么就行了
这边直接给出CMD命令安装。记得切换到你下载的yolov5的文件夹下后再执行以下命令！
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt在pip install的时候，可能会出现read timeout的情况多重复运行几次就行了。
数据集的收集与制作
首先得有数据，你才能训练和制作模型，这一步你可以寻找开源的数据，也可以自己爬虫，制作数据集是一个繁琐且枯燥的过程，这里需要使用LabelImg工具
这里提供链接省的各位找：LabelImg工具下载 提取码：fmcg
具体步骤就不展开了，大家可以看我的这篇博客：使用labelImg标注图像产生yolov5数据集
但是大家标注具体类别的时候一定要细心，因为这会决定你的预测结果。
这是一篇教学文章，为了让大家快速上手，避免浪费大量时间收集数据
数据已经给各位准备好了
安全帽数据集 提取码：03g2 https://pan.baidu.com/s/16UDG3nlx8ZY4QZ4Gmzxc3g%C2%A0数据集是放在百度网盘的（挺大的），有需要的可以在评论区留下你的邮箱，我私发给你！
文件配置
配置训练集测试集划分txt文件
首先在下载好的文件解压到data文件夹内
（或者不下载数据集，直接在data文件夹下创建VOC2028文件夹，文件夹下包含这两个文件）
（注意： JPEGimages内为数据集原始图片，Annotations内为标注的xml文件）
你需要将标注完之后将图片和标注文件放入这两个文件夹内。
分割测试集训练集代码：
import argparse
import os
import random
parser = argparse.ArgumentParser()
# xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='VOC2028/Annotations', type=str, help='input xml label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='VOC2028/ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 0.8 # 训练集和验证集所占比例。 这里划分0.8的测试集，可自己进行调整
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()以上代码存放在data文件夹下（就是和VOC2028文件夹放在一块），我代码叫NO1.py,这个随便自己取都行，叫123.py也可以（这里最好按照命名规则来，我这里是为了方便查找）
执行完毕，查看是否正确执行
数据集内的 VOC2028\ImageSets 文件夹下有个Main文件，这是配置好的训练集和测试集
路径和我一样的就会在 VOC2028\ImageSets 文件夹下有个Main文件，文件内就包含数据集测试集的划分内容
有了这些就可以让程序去划分图像，分为测试集验证集
将划分完的测试训练集xml文件转为yolov5可识别的数据集
NO2.py (路径和NO1.py文件是一样的存放位置，都是放在data文件夹内)
import os
import shutil
from pathlib import Path
from shutil import copyfile
from xml.dom.minidom import parse
import numpy as np
from PIL import Image, ImageDraw
from tqdm import tqdm
my_class = ['helmet', 'person', 'hat'] # xml打上的标签类别
image = '.jpg' # 图像文件的统一格式
FILE_ROOT = Path(r"../data") # data的相对路径
# 原始数据集
IMAGE_SET_ROOT = FILE_ROOT.joinpath(r"VOC2028/ImageSets/Main") # 图片区分文件的路径
IMAGE_PATH = FILE_ROOT.joinpath(r"VOC2028/JPEGImages") # 图片的位置
ANNOTATIONS_PATH = FILE_ROOT.joinpath(r"VOC2028/Annotations") # 数据集标签文件的位置
LABELS_ROOT = FILE_ROOT.joinpath(r"VOC2028/Labels") # 进行归一化之后的标签位置
# YOLO 需要的数据集形式的新数据集
DEST_IMAGES_PATH = Path(r"Safety_Helmet_Train_dataset/score/images") # 区分训练集、测试集、验证集的图片目标路径
DEST_LABELS_PATH = Path(r"Safety_Helmet_Train_dataset/score/labels") # 区分训练集、测试集、验证集的标签文件目标路径
def cord_converter(size, box):
"""
将标注的 xml 文件标注转换为 darknet 形的坐标
:param size: 图片的尺寸： [w,h]
:param box: anchor box 的坐标 [左上角x,左上角y,右下角x,右下角y,]
:return: 转换后的 [x,y,w,h]
"""
x1 = int(box[0])
y1 = int(box[1])
x2 = int(box[2])
y2 = int(box[3])
dw = np.float32(1. / int(size[0]))
dh = np.float32(1. / int(size[1]))
w = x2 - x1
h = y2 - y1
x = x1 + (w / 2)
y = y1 + (h / 2)
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return [x, y, w, h]
def save_label_file(img_jpg_file_name, size, img_box):
"""
保存标签的解析文件
:param img_jpg_file_name:
:param size:
:param img_box:
:return:
"""
save_file_name = LABELS_ROOT.joinpath(img_jpg_file_name).with_suffix('.txt')
with open(save_file_name, "a+") as f:
for box in img_box:
cls_num = my_class.index(box[0])
# print(box[0], cls_num)
"""# 原版代码是需要手动修改类别，改为上述的自动识别了
for box in img_box:
if box[0] == 'person': # 数据集 xml 中的 hat 指的是头
cls_num = 1
elif box[0] == 'hat':
cls_num = 2
else:
continue"""
new_box = cord_converter(size, box[1:]) # 转换坐标
f.write(f"{cls_num} {new_box[0]} {new_box[1]} {new_box[2]} {new_box[3]}\n")
def test_dataset_box_feature(file_name, point_array):
"""
使用样本数据测试数据集的建议框
:param file_name: 图片文件名
:param point_array: 全部的点 [建议框sx1,sy1,sx2,sy2]
:return: None
"""
im = Image.open(IMAGE_PATH.joinpath(file_name).with_suffix(image))
im_draw = ImageDraw.Draw(im)
for box in point_array:
x1 = box[1]
y1 = box[2]
x2 = box[3]
y2 = box[4]
im_draw.rectangle((x1, y1, x2, y2), outline='red')
im.show()
def get_xml_data(img_xml_file: Path):
"""
获取 xml 数据
:param img_xml_file: 图片路径
:return:
"""
dom = parse(str(img_xml_file))
xml_root = dom.documentElement
img_name = xml_root.getElementsByTagName("filename")[0].childNodes[0].data
img_size = xml_root.getElementsByTagName("size")[0]
objects = xml_root.getElementsByTagName("object")
img_w = img_size.getElementsByTagName("width")[0].childNodes[0].data
img_h = img_size.getElementsByTagName("height")[0].childNodes[0].data
img_c = img_size.getElementsByTagName("depth")[0].childNodes[0].data
img_box = []
for box in objects:
cls_name = box.getElementsByTagName("name")[0].childNodes[0].data
x1 = int(box.getElementsByTagName("xmin")[0].childNodes[0].data)
y1 = int(box.getElementsByTagName("ymin")[0].childNodes[0].data)
x2 = int(box.getElementsByTagName("xmax")[0].childNodes[0].data)
y2 = int(box.getElementsByTagName("ymax")[0].childNodes[0].data)
img_box.append([cls_name, x1, y1, x2, y2])
# test_dataset_box_feature(img_xml_file.name, img_box)
save_label_file(img_xml_file.name, [img_w, img_h], img_box)
def copy_data(img_set_source, img_labels_root, imgs_source, dataset_type):
"""
将标签文件和图片复制到最终数据集文件夹中
:param img_set_source: 原数据集图片总路径
:param img_labels_root: 生成的 txt 标签总路径
:param imgs_source:
:param dataset_type: 生成数据集的种类
:return:
"""
file_name = img_set_source.joinpath(dataset_type).with_suffix(".txt") # 获取对应数据集种类的图片
# 判断目标图片文件夹和标签文件夹是否存在，不存在则创建
os.makedirs(FILE_ROOT.joinpath(DEST_IMAGES_PATH, dataset_type), exist_ok=True)
os.makedirs(FILE_ROOT.joinpath(DEST_LABELS_PATH, dataset_type), exist_ok=True)
with open(file_name, encoding="UTF-8") as f:
for img_name in tqdm(f.read().splitlines()):
img_sor_file = imgs_source.joinpath(img_name).with_suffix(image)
label_sor_file = img_labels_root.joinpath(img_name).with_suffix('.txt')
# 复制图片
dict_file = FILE_ROOT.joinpath(DEST_IMAGES_PATH, dataset_type, img_name).with_suffix(image)
copyfile(img_sor_file, dict_file)
# 复制 label
dict_file = FILE_ROOT.joinpath(DEST_LABELS_PATH, dataset_type, img_name).with_suffix('.txt')
copyfile(label_sor_file, dict_file)
if __name__ == '__main__':
root = ANNOTATIONS_PATH # 数据集 xml 标签的位置
if LABELS_ROOT.exists():
# 清空标签文件夹
print("Cleaning Label dir for safety generating label, pls wait...")
shutil.rmtree(LABELS_ROOT)
print("Cleaning Label dir done!")
LABELS_ROOT.mkdir(exist_ok=True) # 建立 Label 文件夹
# 生成标签
print("Generating Label files...")
with tqdm(total=len(os.listdir(root))) as p_bar:
for file in root.iterdir():
p_bar.update(1)
get_xml_data(file)
# 将文件进行 train、val、test 的区分
for dataset_input_type in ["train", "val", "test"]:
# print(f"Copying data {dataset_input_type}, pls wait...")
copy_data(IMAGE_SET_ROOT, LABELS_ROOT, IMAGE_PATH, dataset_input_type)
运行之后和我一样，那么恭喜你完成一半了 ！
等到百分之百之后就完成了，会在文件下生成 Safety_Helmet_Train_dataset 文件
一定要检查一下 yolov5-master\data\Safety_Helmet_Train_dataset\Labels 文件夹下的txt 文件是否拥有数据！！！！
没有数据就代表你之前设置的类别有问题
解决方法：上一步中 NO2.py 文件内classes 数组的类别重新修改为自己的类别，必须与xml标注文件内的标签一致！
正确操作步骤之后 执行的效果如下：
到这里你已经完成了训练模型的所有文件配置。可以开始准备着手训练模型了！
创建自己的模型
下载预训练模型
进入到 yolov5-master\weights 文件夹放置预训练模型，这里我使用的是small版本的模型，模型更小训练更快
因为链接老是被和谐，这里给了github的官方下载地址。 = =
11-25又又又更新了一次，再被和谐只能去GItHub搜索yolov5在这点击下载就可以了
（这里使用的是轻量部署模型，yolov5s，训练速度快，需要更好的模型可以在官方网站下载）
然后放置到
在路径 yolov5-master/data/ 文件夹中创建自己的数据集配置文件custom_data.yaml：
代码：
# Custom data for safety helmet
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
train: ./data/Safety_Helmet_Train_dataset/score/images/train
val: ./data/Safety_Helmet_Train_dataset/score/images/val
nc: 3 # number of classes
names: ['person', 'head', 'helmet'] # my classes训练模型
在cmd中进入yolov5 文件夹（在编译器内进入也行，我自己是在Pycharm内进入的，效果一样）
方法一、
方法二、
我这里是在Pycharm中打开终端，没有Pycharm也可以选cmd打开（效果一样）
记得进去yolov5的目录下
然后输入以下指令：
python train.py --img 640 --batch 16 --epochs 50 --data ./data/custom_data.yaml --weights ./weights/yolov5s.pt注意
--batch 16 一次喂多少数据，我这内存就能给16，所以可以不传按默认16
（如果内存比较小的，建议改为8，或者4）
（但是也需要根据实际情况，看自身内存的）
--epochs 50 代表迭代五十次，可自行修改
--data ./data/custom_data.yaml 代表数据集配置文件
--weights 代表预训练模型
当你等待一段时间之后，看见以下图片，就证明模型训练成功了！在你等待大概十七个小时之后就可以拥有自己的数据模型PT了！（GPU训练很快，下次出个GPU安装的教程）
（因为我没有使用GPU进行训练，在运行的时候一直提示我CUDA不可用，一般忽略就行。 ）
warnings.warn('User provided device_type of \'cuda\', but CUDA is not available. Disabling'）
大概意思就是我可以使用GPU进行训练，但是我没有使用
模型保存位置： yolov5-master\runs\train\exp2\weights\best.pt
训练完的模型怎么使用呢？
更新：2022年 02月02日 14：37
更新：2022年 07月01日 15：57
1309
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
