yolov5数据构建的两种方式

  对于还不太熟悉yolov5源代码的同学，利用自己的数据集构建符合yolov5训练格式的数据集应该挺让人头大的，所以下面我来讲一下构建训练数据的两种常用方式：
  一般标注的标签格式为xml文件格式，我们首先需要将xml格式转成txt格式文件。

import os
import xml.etree.ElementTree as ET
import json
import yaml
def yaml_load(file='data.yaml'):
    # Single-line safe yaml loading
    with open(file, errors='ignore') as f:
        return yaml.safe_load(f)

def get_and_check(root, name, length):
    vars = root.findall(name)
    if len(vars) == 0:
        raise NotImplementedError('Can not find %s in %s.'%(name, root.tag))
    if length > 0 and len(vars) != length:
        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.'%(name, length, len(vars)))
    if length == 1:
        vars = vars[0]
    return vars

def voc_to_coco(xml_path, save_path, classes_dict):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    size = get_and_check(root, 'size', 1)
    img_width = int(get_and_check(size, 'width', 1).text)
    img_height = int(get_and_check(size, 'height', 1).text)
    f = open(save_path + "/" + os.path.basename(xml_path).replace('.xml', '.txt'), 'w')
    for index,  obj in enumerate(root.findall('object')):
        cls = obj.find('name').text
        category_id = next(k for k, v in classes_dict.items() if v == cls)
        xmin = int(obj.find('bndbox').find('xmin').text)
        ymin = int(obj.find('bndbox').find('ymin').text)
        xmax = int(obj.find('bndbox').find('xmax').text)
        ymax = int(obj.find('bndbox').find('ymax').text)
        assert(xmax > xmin), "xmax <= xmin, {}".format(xml_path)
        assert(ymax > ymin), "ymax <= ymin, {}".format(xml_path)
        o_width = abs(xmax - xmin)
        o_height = abs(ymax - ymin)
        xcenter = xmin + o_width / 2
        ycenter = ymin + o_height / 2
        try:
            xcenter = round(xcenter / img_width, 6)
            ycenter = round(ycenter / img_height, 6)
        except:
            print(xml_path)
        o_width = round(o_width / img_width, 6)
        o_height = round(o_height / img_height, 6)
        info = [str(i) for i in [category_id, xcenter, ycenter, o_width, o_height]]

        if index == 0:
            f.write(" ".join(info))
        else:
            f.write("\n" + " ".join(info))

if __name__ == '__main__':
	# 这两行代码是为了产生类别标签对应。 如{0："person", 1:"cat"....}如果类别少可以自己手写
	# 我这里通过读取构建自己数据集的yaml文件，里面的类别确定了，这个文件就像coco128.yaml文件一样
	# xml转txt不是本文重点，当然还有其他更方便的办法，请自行查阅
    butt_yaml_path = "yolov5-7.0/data/butterfly.yaml"
    classes_dict = yaml_load(butt_yaml_path)['names']
    Annotation_path = "Annotations"  这里是xml文件路径
    save_txt_path = "lables" 这里是要保存txt文件地址，自己定义
    anno_list = os.listdir(Annotation_path)
    for anno_name in anno_list:
        per_anno_path = os.path.join(Annotation_path, anno_name)
        voc_to_coco(per_anno_path, save_txt_path, classes_dict)

  通过以上代码就可以获取labels文件夹，里面是每张图片的GT框信息。

  还有images文件夹，里面是对应的图像文件。

方式1

import os
import random
import shutil

images_path = "images" #图像路径
# 这里几个文件夹可以手动创建好
root_dir = os.path.dirname(images_path)
save_train_image_path = os.path.join(root_dir, "train", "images")
save_train_label_path = os.path.join(root_dir, "train", "labels")

save_val_image_path = os.path.join(root_dir, "val", "images")
save_val_label_path = os.path.join(root_dir, "val", "labels")


if not os.path.exists(save_train_image_path):
    os.makedirs(save_train_image_path)
if not os.path.exists(save_train_label_path):
    os.makedirs(save_train_label_path)

if not os.path.exists(save_val_image_path):
    os.makedirs(save_val_image_path)
if not os.path.exists(save_val_label_path):
    os.makedirs(save_val_label_path)

# 我这里的images文件夹下就一层 要是有多层可以用os.walk()
train_rate = 0.8
imgs_name_list = os.listdir(images_path)
num = len(imgs_name_list)
list_index = range(num)
tr = int(num * train_rate)
train = random.sample(list_index, tr)

for i in list_index:
    image_path = os.path.join(images_path, imgs_name_list[i])
    sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}'  # /images/, /labels/ substrings
    label_path = sb.join(image_path.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt'
    if i in train:
        shutil.copy(image_path, save_train_image_path)
        shutil.copy(label_path, save_train_label_path)
    else:
        shutil.copy(image_path, save_val_image_path)
        shutil.copy(label_path, save_val_label_path)

  coco128数据集就是这样构建的，这应该是很多同学的构建方式，但是这种方式会将数据复制一遍多占内存，就算直接move，这样就把原数据改动了，感觉不太好。

方式2

  通过写train.txt文件和val.txt文件，这种方式不会改动原图像文件夹，只需要将图像路径（最好记录绝对路径）记录在txt文件中，这种方式比较推荐。（把图像的绝对路径记录就好了）

# coding:utf-8

import os
import random
import argparse
parser = argparse.ArgumentParser()
#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--image_path', default='images', type=str, help='input xml label path')
#保存自己构建的train.txt和val.txt的路径
parser.add_argument('--txt_save_path', default='', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.8
imgfilepath = opt.image_path
txtsavepath = opt.txt_save_path
total_img = os.listdir(imgfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_img)
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:
    img_path = os.path.join(imgfilepath, total_img[i])
    name = img_path + '\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()

  在构建数据集yaml文件时，需要修改数据途径如下

  不知道各位同学喜欢用那种方式，欢迎沟通讨论。。。。。。。。

注意点：
  图像文件夹和标签文件夹要用 images和labels命名，因为在源码中是通过这两个名字找图像和标签的