数据集转换-VOC2YOLOv5

一键转换，开箱即用

需要设置的参数只有2个，就是Annotations和images，其他的使用默认参数即可，假如没有需要的话，再修改一下比例参数或者其他的。

def xml2voc(xmlfilepath = 'Annotations',trainval_percent = 0.8,train_percent = 0.8):
    def makedir(path):
        import os
        try:
            os.makedirs(path)
        except:
            pass

    import os
    import random
    # trainval_percent = 0.8          #训练验证集占所有数据集的比重
    # train_percent = 0.8             #训练集占训练验证集的比重
    # xmlfilepath = 'Annotations'         ##xml文件所在的文件夹
    total_xml = os.listdir(xmlfilepath)

    num = len(total_xml)
    list = range(num)
    tv = int(num * trainval_percent)
    tr = int(tv * train_percent)
    trainval = random.sample(list, tv)
    train = random.sample(trainval, tr)
    makedir('ImageSets/Main')

    # filehandle = open("ImageSets/Main/trainval.txt", "w")
    # filehandle.close()
    ftrainval = open('ImageSets/Main/trainval.txt', 'w')
    ftest = open('ImageSets/Main/test.txt', 'w')
    ftrain = open('ImageSets/Main/train.txt', 'w')
    fval = open('ImageSets/Main/val.txt', 'w')

    for i in list:
        name = total_xml[i][:-4] + '\n'
        if i in trainval:
            ftrainval.write(name)
            if i in train:
                ftrain.write(name)
            else:
                fval.write(name)
        else:
            ftest.write(name)

    ftrainval.close()
    ftrain.close()
    fval.close()
    ftest.close()
def VOC_YOLO_txt(classes):
    import xml.etree.ElementTree as ET
    import pickle
    import os
    from os import listdir, getcwd
    from os.path import join
    sets = ['train', 'test', 'val']
    # classes = ['fruit']  # 需要更改

    def convert(size, box):
        dw = 1. / size[0]

        dh = 1. / size[1]
        x = (box[0] + box[1]) / 2.0
        y = (box[2] + box[3]) / 2.0
        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('Annotations/%s.xml' % (image_id))  # 需要更改
        out_file = open('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
            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')

    wd = getcwd()
    print(wd)
    for image_set in sets:
        if not os.path.exists('labels/'):
            os.makedirs('labels/')
        image_ids = open('ImageSets/Main/%s.txt' % (image_set)).read().strip().split()  # 需要更改
        list_file = open('%s_sets_directory.txt' % (image_set), 'w')
        labels_list_file = open('%s_labels_directory.txt' % (image_set), 'w')
        for image_id in image_ids:
            list_file.write('images/%s.jpg\n' % (image_id))
            labels_list_file.write('labels/%s.txt\n' % (image_id))
            convert_annotation(image_id)
        list_file.close()
        labels_list_file.close()

def txt_2_yolo_v5(imageset = 'fruit'):
    def makedir(path):
        import os
        try:
            os.makedirs(path)
        except:
            pass

    def get_picture_labels(object_directry, target_directoty):
        import os
        import shutil
        with open(object_directry) as  f:  # mark1,根据输出的地址调节
            name_result = f.readlines()
            for i, str in enumerate(name_result):
                name_result[i] = (str.replace('\n', ''))
                # print('now  i is {},str is {}'.format(i , str)) #用于显示文件名或调试
                shutil.copy(name_result[i], target_directoty)
        print((name_result))


    makedir('./{}/train/images'.format(imageset))
    makedir('./{}/train/labels'.format(imageset))
    makedir('./{}/val/images'.format(imageset))
    makedir('./{}/val/labels'.format(imageset))
    makedir('./{}/test/images'.format(imageset))
    makedir('./{}/test/labels'.format(imageset))
    object_directry = './train_sets_directory.txt'
    img_target_directoty = './{}/train/images'.format(imageset)
    get_picture_labels(object_directry, img_target_directoty)



    object_directry = './train_labels_directory.txt'
    img_target_directoty = './{}/train/labels'.format(imageset)
    get_picture_labels(object_directry, img_target_directoty)


    object_directry = './test_sets_directory.txt'
    img_target_directoty = './{}/test/images'.format(imageset)
    get_picture_labels(object_directry, img_target_directoty)
    '''
    #transfer valid label
     '''
    object_directry = './test_labels_directory.txt'
    img_target_directoty = './{}/test/labels'.format(imageset)
    get_picture_labels(object_directry, img_target_directoty)

    object_directry = './val_labels_directory.txt'
    img_target_directoty = './{}/val/images'.format(imageset)
    get_picture_labels(object_directry, img_target_directoty)

    object_directry = './val_labels_directory.txt'
    img_target_directoty = './{}/val/labels'.format(imageset)
    get_picture_labels(object_directry, img_target_directoty)


xml2voc(xmlfilepath = 'Annotations',trainval_percent = 0.8,train_percent = 0.8)

classes= ['fruit']
VOC_YOLO_txt(classes)

txt_2_yolo_v5()

生成的目录树如下