1. 将数据集标注成yolo格式
这里我曾经已经写过一个对xml目录文件转txt标注文件的脚本，修改目录路径直接运行即可。脚本文件见博文：yolov5的数据集准备 | 处理Pascal voc格式的数据集
xml_to_txt.py
import os import pickle import xml.etree.ElementTree as ET from tqdm import tqdm # chance your classes here names: ['with_mask', 'without_mask', 'mask_weared_incorrect'] # function: # (xmin, xmax, ymin, ymax) -> (center_x, cneter_y, box_weight, box_height) # size: (w, h) # box : (xmin, xmax, ymin, ymax) def convert(size, box): dw = 1.0 / size[0] dh = 1.0 / size[1] x = (box[0] + box[1]) / 2.0 # center_x y = (box[2] + box[3]) / 2.0 # center_y w = box[1] - box[0] # box_weight h = box[3] - box[2] # box_height x = x * dw w = w * dw y = y * dh h = h * dh return (x, y, w, h) # 功能：对单个文件进行转换，其中名字是相匹配的 # infile: xml文件路径 # outfile：txt文件输出路径 def convert_annotation(infile, outfile): in_file = open(infile, encoding='utf-8') # xml path out_file = open(outfile, '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') in_file.close() out_file.close() # 功能：对xml目录下的全部xml文件全部转换成yolo的txt格式，并输出在savepath目下 # xmlpath: the dir path save xml files # savepath: the dir path to save txt files def xml_to_txt(xmlpath, savepath): assert os.path.exists(xmlpath), "{} not exists.".format(xmlpath) if not os.path.exists(savepath): os.mkdir(savepath) assert os.path.exists(savepath), "{} create fail.".format(savepath) xmlfiles = os.listdir(xmlpath) assert len(xmlfiles) != 0, "{} have no files".format(xmlpath) xmlfilepaths = [os.path.join(xmlpath, xmlfile) for xmlfile in xmlfiles] for xmlfilepath in tqdm(xmlfilepaths, desc="change xml to txt: "): # E:\学习\机器学习\数据集\众盈比赛数据集\xml\train1.xml -> train1 image_id = xmlfilepath.split('\\')[-1].split('.')[0] # train1 -> \..\train1.txt txtfilepath = os.path.join(savepath, image_id) + ".txt" convert_annotation(xmlfilepath, txtfilepath) if __name__ == '__main__': xmlpath = r"E:\学习\机器学习\数据集\VOC2012\VOCdevkit\VOC2012\Annotations" savepath = r"E:\学习\机器学习\数据集\VOC2012\VOCdevkit\VOC2012\YoloLabels" xml_to_txt(xmlpath, savepath)
2. 生成train.txt与val.txt文件
最重要的就是这一步骤了，需要分别保存所有的训练与测试图片路径在txt文件中。
那么，首先就是需要构建验证集与测试集，然后再根据对应的数据集来存放相应的图片路径。
2.1 数据集结构处理
对于这里，yolov3的数据集结构与yolov5的不一致，所以需要重新摆放数据集结构，结构如下：
├── Dataset 自定义数据集根目录 │ ├── train 训练集目录 │ │ ├── images 训练集图像目录 │ │ └── labels 训练集标签目录 │ └── val 验证集目录 │ ├── images 验证集图像目录 │ └── labels 验证集标签目录
对于之前的代码进行简单的修改即可（就是主要是修改路径）
构建数据集结构代码如下：
from sklearn.model_selection import train_test_split import os from tqdm import tqdm from shutil import copyfile # 功能: 递归穿件目录文件，有则跳过，没有则创建并验证 def _mkdir(dirpath): if not os.path.exists(dirpath): os.makedirs(dirpath) assert os.path.exists(dirpath), "{} create fail.".format(dirpath) # 功能: 从图片路径中获取最后的文件名(含后缀或者不含后缀) def get_filename(filepath, suffix=True): if suffix is True: image_id = filepath.split('\\')[-1] # 获取文件名(含后缀) else: image_id = filepath.split('\\')[-1].split('.')[0] # 获取文件名(不含后缀) return image_id # 功能: 根据图片数据与标签数据构建yolov5格式的数据集 # imagepath: 存放数据图片的路径 # labelpath: 存放数据标签的路径 # datasetpath: 构建voc格式数据集的目录路径 def dataset_split(imagepath, labelpath, datasetpath): assert os.path.exists(imagepath), "{} not exists.".format(imagepath) assert os.path.exists(labelpath), "{} not exists.".format(labelpath) # 构建图片与标签的路径列表，然后按8:2进行训练与测试的切分 images = os.listdir(imagepath) imagelists = [os.path.join(imagepath, image) for image in images] assert len(imagelists) != 0, "{} have no files".format(imagepath) labels = os.listdir(labelpath) labellists = [os.path.join(labelpath, label) for label in labels] assert len(labellists) != 0, "{} have no files".format(labelpath) X_train, X_test, y_train, y_test = train_test_split(imagelists, labellists, test_size=0.2, random_state=42) # 构建voc格式的数据集存放格式 _mkdir(datasetpath) # 在数据路径datasetpath下，分别递归建立目录: # | -----Dataset # | -----train # | -----images # | -----labels # | -----val # | -----images # | -----labels images_train_dirpath = os.path.join(datasetpath, "train", "images") _mkdir(images_train_dirpath) labels_train_dirpath = os.path.join(datasetpath, "train", "labels") _mkdir(labels_train_dirpath) images_val_dirpath = os.path.join(datasetpath, "val", "images") _mkdir(images_val_dirpath) labels_val_dirpath = os.path.join(datasetpath, "val", "labels") _mkdir(labels_val_dirpath) # 复制图片/标签列表中的数据到指定子目录下 for image_path in tqdm(X_train, desc="process train images"): image_newpath = os.path.join(images_train_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) for image_path in tqdm(X_test, desc="process val images"): image_newpath = os.path.join(images_val_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) for image_path in tqdm(y_train, desc="process train labels"): image_newpath = os.path.join(labels_train_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) for image_path in tqdm(y_test, desc="process val labels"): image_newpath = os.path.join(labels_val_dirpath, get_filename(image_path)) copyfile(image_path, image_newpath) print("Process Success. The data has save to {}".format(datasetpath)) if __name__ == '__main__': imagepath = r"E:\学习\机器学习\数据集\mask\images" labelpath = r"E:\学习\机器学习\数据集\mask\YoloLabels" datasetpath = r"E:\学习\机器学习\数据集\mask\Yolov3_Dataset" dataset_split(imagepath, labelpath, datasetpath)
2.2 文件名保存处理
接着，先遍历刚刚整理好的数据集结构，然后把图像全部路径的放在对应的txt文件中。
文件名保存代码如下：
import os import numpy as np def make_filepath(dirpath, savepath): assert os.path.exists(dirpath), "{} not exists.".format(dirpath) assert os.path.exists(savepath), "{} not exists.".format(savepath) trainpath = os.path.join(dirpath, "train", "images") validpath = os.path.join(dirpath, "val", "images") assert os.path.exists(trainpath), "{} not exists.".format(trainpath) assert os.path.exists(validpath), "{} not exists.".format(validpath) trainfiles = os.listdir(trainpath) assert len(trainfiles) > 0, "{} has no files.".format(trainpath) validfiles = os.listdir(validpath) assert len(validfiles) > 0, "{} has no files.".format(validpath) with open(os.path.join(savepath, "train.txt"), "w") as f: for trainfile in trainfiles: trainfile = os.path.join(r"./dataset/mask/train/images", trainfile) + "\n" f.writelines(trainfile) with open(os.path.join(savepath, "val.txt"), "w") as f: for validfile in validfiles: validfile = os.path.join(r"./dataset/mask/val/images", validfile) + "\n" f.writelines(validfile) print("write success...") if __name__ == '__main__': # dirpath = r"E:\PyCharm\workspace\Detection\yolov3_spp\dataset\mask" # savepath = r"E:\PyCharm\workspace\Detection\yolov3_spp\data" imgpath = r"./dataset/mask" savepath = r"./data" make_filepath(imgpath, savepath)
处理好之后，会生成两个txt文件，分别对应的是训练集的全部图像路径和验证集图像的全部路径。可以注意到，其实这里是不需要处理label路径的，因为会根据图像的路径自动的找到label的路径。处理之后的结果如下所示：
train.txt
val.txt
3. 创建一个.names标签
.names标签文件用来存储类别信息
这里我的口罩检测数据集是有3类：[‘with_mask’, ‘without_mask’, ‘mask_weared_incorrect’]，然后直接记录在文件中即可。
4. 创建一个.data文件
data文件包含txt文件的路径与标签文件的路径，还需要指定类别数量
这里根据项目的路径位置自行设置：
5. 更改预测器卷积核的个数
由于这里我的类别个数为3，所以[convolutional]中的filters值为：（5+3）*3=24，然后修改一些[yolo]中的classes即可，这里改为你的类别数，如下所示：
6. 更改超参及训练
这里的超参没有一点炼丹的经验一般不建议改动
然后，就可以开始训练自己的数据集了。
补充一下我对yolov3spp数据集制作的思路：
可以先将数据集按yolov3的结构啦进行摆放，然后再生成相对应的文件，毕竟需要使用不同的系统（服务器与本地），这样做可以不需要改动摆放结构，只需要改变相关文件的路径即可。所以大体思路如下，也就是我上面所提供的两个脚本
1）制作对应结构的数据集；
2）根据生成的数据集来生成相关文件，比如训练图片路径与测试图片路径；
