查看yolov5的配置文件coco.yaml和coco128.yaml,yolov5数据集可以有两种形式,
一种像coco128所示,在images和labels文件夹下,建立train、test、val三个文件夹,存放训练集、测试集、验证集的标签和数据,然后再复制coco128.yaml前面路径部分,改为自己的路径,其中,训练集、测试集、验证集均写相应的文件夹路径,再将类别数量、名称改为自己数据集的名称和数量。download部分完全没用,不要复制
另一种就是将训练集、测试集、验证集图片路径直接写入txt文件中,复制coco.yaml,路径修改为自己的txt文件、种类数量、名称都要改为自己的数据
VOC数据集如图所示:Annotations_original是数据集没有的,是我调整数据后加的,JPEGImages下存放的是图片,Annotations存放标签,ImageSets下的main文件夹存放train/test/val文件标签
1.JPEGImages文件夹名称需要改为images(手动修改)
2.需要存放训练集、测试集、验证集图片路径的txt文件
3.labels需要由xml文件改为yolo格式
运行下面的代码即可:
- import os
- import random
- import xml.etree.ElementTree as ET
-
- import numpy as np
-
- # --------------------------------------------------------------------------------------------------------------------------------#
- # annotation_mode用于指定该文件运行时计算的内容
- # annotation_mode为0代表整个标签处理过程,训练集、测试集、验证集图片路径的txt文件以及标签xml文件改为yolo格式
- # annotation_mode为1代表获得训练集、测试集、验证集图片路径的txt文件
- # annotation_mode为2代表标签xml文件改为yolo格式
- # --------------------------------------------------------------------------------------------------------------------------------#
- annotation_mode = 0
- # -------------------------------------------------------------------#
- # 必须要修改,用于标签的目标信息
- # 与训练和预测所用的classes_path一致即可
- # 如果生成的2007_train.txt里面没有目标信息
- # 那么就是因为classes没有设定正确
- # 仅在annotation_mode为0和2的时候有效
- # -------------------------------------------------------------------#
- classes_path = '../Detection/classes.txt'
- # --------------------------------------------------------------------------------------------------------------------------------#
- # trainval_percent用于指定(训练集+验证集)与测试集的比例,默认情况下 (训练集+验证集):测试集 = 9:1
- # train_percent用于指定(训练集+验证集)中训练集与验证集的比例,默认情况下 训练集:验证集 = 9:1
- # 仅在annotation_mode为0和1的时候有效
- # --------------------------------------------------------------------------------------------------------------------------------#
- trainval_percent = 0.8
- train_percent = 0.8
- # -------------------------------------------------------#
- # 指向VOC数据集所在的文件夹
- # 默认指向根目录下的VOC数据集
- # -------------------------------------------------------#
- VOCdevkit_path = '../Detection'
-
-
- def get_classes(classes_path):
- with open(classes_path, encoding='utf-8') as f:
- class_names = f.readlines()
- class_names = [c.strip() for c in class_names]
- return class_names, len(class_names)
-
-
- classes, _ = get_classes(classes_path)
- # -------------------------------------------------------#
- # 统计目标数量
- # -------------------------------------------------------#
- photo_nums = np.zeros(len(VOCdevkit_sets))
- nums = np.zeros(len(classes))
-
-
- def convert_label(path, image_id):
- def convert_box(size, box):
- dw, dh = 1. / size[0], 1. / size[1]
- x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
- return x * dw, y * dh, w * dw, h * dh
-
- in_file = open(os.path.join(path,f'VOC2007/Annotations/{image_id}.xml'),encoding='utf-8')
- out_file = open(os.path.join(path, f'VOC2007/labels/{image_id}.txt'), 'w',encoding='utf-8')
- 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'):
- cls = obj.find('name').text
- if cls in classes and not int(obj.find('difficult').text) == 1:
- xmlbox = obj.find('bndbox')
- bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
- cls_id = classes.index(cls) # class id
- out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')
-
-
- if __name__ == "__main__":
- random.seed(42)
- if " " in os.path.abspath(VOCdevkit_path):
- raise ValueError("数据集存放的文件夹路径与图片名称中不可以存在空格,否则会影响正常的模型训练,请注意修改。")
-
- if annotation_mode == 0 or annotation_mode == 1:
- print("Generate train/test/val txt")
- xmlfilepath = os.path.join(VOCdevkit_path, 'VOC2007/Annotations')
- saveBasePath = os.path.join(VOCdevkit_path, 'VOC2007')
- temp_xml = os.listdir(xmlfilepath)
- total_xml = []
- for xml in temp_xml:
- if xml.endswith(".xml"):
- total_xml.append(xml)
-
- num = len(total_xml)
- list = list(range(num))
- tv = int(num * trainval_percent)
- tr = int(tv * train_percent)
- trainval = random.sample(list, tv)
- train = random.sample(trainval, tr)
- random.shuffle(list)
-
- print("train and val size", tv)
- print("train size", tr)
- ftrainval = open(os.path.join(saveBasePath, 'trainval.txt'), 'w',encoding='utf-8')
- ftest = open(os.path.join(saveBasePath, 'test.txt'), 'w',encoding='utf-8')
- ftrain = open(os.path.join(saveBasePath, 'train.txt'), 'w',encoding='utf-8')
- fval = open(os.path.join(saveBasePath, 'val.txt'), 'w',encoding='utf-8')
-
- for i in list:
- name = total_xml[i][:-4]
- if i in trainval:
- ftrainval.write('%s/VOC2007/images/%s.jpg\n'% (os.path.abspath(VOCdevkit_path), name))
- if i in train:
- ftrain.write('%s/VOC2007/images/%s.jpg\n'% (os.path.abspath(VOCdevkit_path), name))
- else:
- fval.write('%s/VOC2007/images/%s.jpg\n'% (os.path.abspath(VOCdevkit_path), name))
- else:
- ftest.write('%s/VOC2007/images/%s.jpg\n'% (os.path.abspath(VOCdevkit_path), name))
-
- ftrainval.close()
- ftrain.close()
- fval.close()
- ftest.close()
- print("Generate train/test/val.txt done.")
-
- if annotation_mode == 0 or annotation_mode == 2:
- print("Generate coco labels")
- if not os.path.exists(os.path.join(VOCdevkit_path,'VOC2007/labels/')):
- os.makedirs(os.path.join(VOCdevkit_path,'VOC2007/labels/'))
- image_ids = os.listdir(os.path.join(VOCdevkit_path, 'VOC2007/Annotations'))
- for image_id in image_ids:
- convert_label(VOCdevkit_path, image_id[:-4])
- print("Generate coco labels done")
最终得到的文件情况:
最后按照第一点说的,复制 配置文件并修改可以了
data指向自己的yaml文件,训练权重、epochs等按照自己的情况修改