VOC格式数据集转YOLO格式(xml转txt)

voc格式数据集

yolov5训练需要的格式

主要操作
1.按照voc下的train.txt和val.txt将图片放到images文件夹下的train2017和val2017文件夹下。
2.将对应的xml文件转为txt文件，并放到labels文件夹下的train2017和val2017文件夹下。

代码如下

import os
from PIL import Image
import xml.etree.ElementTree as ET
'''
将voc格式中的xml文件转换为支持yolo训练的txt文件
'''
#设置好自己数据集的类别
# class_list =  ['object', 'animal']  # class names
class_list = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
        'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']  # class names  # class names

def xymin_xymax2x_y_cx_cy(xmin, ymin, xmax, ymax, width, height):
    x_center = xmin + (xmax - xmin) / 2
    y_center = ymin + (ymax - ymin) / 2
    w = xmax - xmin
    h = ymax - ymin

    #归一化

    x = round(x_center / width, 6)
    y = round(y_center / height, 6)
    w = round(w / width, 6)
    h = round(h / height, 6)

    return x, y, w, h


def class_index(name:str):
    '''
    返回对应类别的索引
    :param name:
    :return:
    '''
    for i, class_ in enumerate(class_list):
        if name == class_:
            return i
def voc2yolo(train_txt_path, val_txt_path, image_dir_path, train_image_save_path, val_image_save_path, xml_dir_path, txt_dir_train_path, txt_dir_val_path):
    '''

    :param train_txt_path: VOC中train.txt文件路径
    :param val_txt_path:   VOC中val.txt文件路径
    :param image_dir_path: VOC数据集下保存图片的文件夹路径
    :param train_image_save_path: 训练图片需要保存的文件夹 images/train2017
    :param val_image_save_path:   测试图片需要保存的文件夹 images/val2017
    :param xml_dir_path:          VOC数据集下xml文件夹路径
    :param txt_dir_train_path:    训练数据集对应的txt的文件路径 labels/train2017
    :param txt_dir_val_path:      测试数据集对应的txt的文件路径 labels/val2017
    :return:
    '''
    #按照train.txt和val.txt将图片放到images文件夹下的train2017和val2017文件夹下
    train_list = []
    with open(train_txt_path,"r") as f:
        for line in f:
            train_list.append(line[:-1])
    # print(train_list)
    #
    val_list = []
    with open(val_txt_path, "r") as f:
        for line in f:
            val_list.append(line[:-1])
    # print(val_list)

    all_images_list = []
    for image in os.listdir(image_dir_path):
        new_image = image.split(".")[0]
        all_images_list.append(new_image)
        img = Image.open(os.path.join(image_dir_path, image))
        if new_image in train_list:
            if not os.path.exists(train_image_save_path):
                os.makedirs(train_image_save_path)
            img.save(os.path.join(train_image_save_path,image))
        else:
            if not os.path.exists(val_image_save_path):
                os.makedirs(val_image_save_path)
            img.save(os.path.join(val_image_save_path,image))
    
    #读取xml文件，将其转为txt文件保存
    #train_label_list
    for xml in os.listdir(xml_dir_path):
        #
        tree = ET.ElementTree(file=os.path.join(xml_dir_path,xml))
        root = tree.getroot()
        size = root.find("size")
        width = float(size.find("width").text)
        height = float(size.find("height").text)
        object = root.findall("object")
        new_xml = xml.split(".")[0]
        if new_xml in train_list:
            save_path = txt_dir_train_path
        else:
            save_path = txt_dir_val_path
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        with open(save_path + "/" + xml.split(".")[0] + ".txt", "a+") as f:
            for child in object:
                name = child.find("name").text
                i = class_index(name)
                bndbox = child.find("bndbox")
                dict_xy ={}
                for j in bndbox:
                    dict_xy[j.tag] = j.text
                xmin = float(dict_xy['xmin'])
                ymin = float(dict_xy['ymin'])
                xmax = float(dict_xy['xmax'])
                ymax = float(dict_xy['ymax'])
                print(i, xmin, ymin, xmax, ymax)
                # 进一步检查数据，有的标注信息中可能有w或h为0的情况，这样的数据会导致计算回归loss为nan
                if xmax <= xmin or ymax <= ymin:
                    print("Warning: in '{}' xml, there are some bbox w/h <=0".format(xml))
                    continue
                x, y, w, h = xymin_xymax2x_y_cx_cy(xmin, ymin, xmax, ymax, width, height)
                info = str(i) + " " + str(x) + " " + str(y) + " " + str(w) + " " + str(h) + "\n"
                f.write(info)


if __name__ == "__main__":

    voc2yolo("C:\\VOCdevkit\\VOC2012\\ImageSets\\Main\\train.txt",
             "C:\\VOCdevkit\\VOC2012\\ImageSets\\Main\\val.txt",
             "C:\\VOCdevkit\\VOC2012\\JPEGImages",
             "C:\\voc_test2coco\\images\\train2017",
             "C:\\voc_test2coco\\images\\val2017",
             "C:\\VOCdevkit\\VOC2012\\Annotations",
             "C:\\voc_test2coco\\labels\\train2017",
             "C:\\voc_test2coco\\labels\\val2017")

总结：
代码有些冗长，没有优化，但是是有效的，路径设置对就直接生成v5要的格式。
ET模块也不是很会用，bndbox那边用字典获取xmin，xmax。。。有空在认真学一下吧。。