目标检测：数据集划分 XML数据集转YOLO标签

文章目录

1、前言：
2、生成对应的类名
3、xml转为yolo的label形式
4、优化代码
5、划分数据集
6、画目录树
7、目标检测系列文章

1、前言：

本文演示如何划分数据集，以及将VOC标注的xml数据转为YOLO标注的txt格式，且生成classes的txt文件。

# 本文演示的项目目录
E:
└──dataset
        ├── images_package  # 存放图片文件夹
        │    ├── 000002.jpg
        │    ├── 000003.jpg
        │    ├── 000004.jpg
        │    ├── 000005.jpg
        │    ├── 000006.jpg
        │    ├── zebra_crossing_20180129-000545_362.jpg
        │    ├── zebra_crossing_20180129-000545_364.jpg
        │    ├── zebra_crossing_20180129-000545_366.jpg
        │    └── zebra_crossing_20180129-000645_368.jpg
        └── xml_outputs  # 存放图片对应的XML
            ├── 000002.xml
            ├── 000003.xml
            ├── 000004.xml
            ├── 000005.xml
            ├── 000006.xml
            ├── zebra_crossing_20180129-000545_362.xml
            ├── zebra_crossing_20180129-000545_364.xml
            ├── zebra_crossing_20180129-000545_366.xml
            └── zebra_crossing_20180129-000645_368.xml

2、生成对应的类名

创建create_classes_json.py自动生成对应的类名json文件，以及在控制台输出对应的类名集。

from doctest import REPORTING_FLAGS
from lib2to3.pgen2.token import RPAR
import os
from tqdm import tqdm
from lxml import etree
import json
 
 
# 读取 xml 文件信息，并返回字典形式
def parse_xml_to_dict(xml):
    if len(xml) == 0:  # 遍历到底层，直接返回 tag对应的信息
        return {xml.tag: xml.text}
 
    result = {}
    for child in xml:
        child_result = parse_xml_to_dict(child)  # 递归遍历标签信息
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:  # 因为object可能有多个，所以需要放入列表里
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}
 
 
# 提取xml中name保留为json文件
def xml2json(data,json_path):
    xml_path = [os.path.join(data, i) for i in os.listdir(data)]
    classes = []      # 目标类别
    num_object = 0
    for xml_file in tqdm(xml_path, desc="loading..."):
        with open(xml_file,encoding='gb18030',errors='ignore') as fid:      # 防止出现非法字符报错
            xml_str = fid.read()
        xml = etree.fromstring(xml_str)
        data = parse_xml_to_dict(xml)["annotation"]  # 读取xml文件信息
        for j in data['object']:        # 获取单个xml文件的目标信息
            ob = j['name']
            num_object +=1
            if ob not in classes:
                classes.append(ob)
    print(num_object)
    # 生成json文件
    labels = {}
    for index,object in enumerate(classes):
        labels[index] = object

    # 打印类名
    classes_name=[labels[key] for key in labels]
    print(f'类名：{classes_name}')
    # 打印类型字典
    print(f'字典形式：{labels}')

    # json.dumps将python对象转为json对象（将dict转化成str）。 json.loads将json字符串解码成python对象（将str转化成dict）
    labels = json.dumps(labels,indent=4)
    json_path=os.path.join(json_path,'classes_indices.json')
    with open(json_path,'w') as f:
        f.write(labels)
 
 
if __name__ == "__main__":
    # 数据集的 xml 目录
    xml_path = 'E:\\dataset\\xml_outputs' 
    # 存放类名json路径
    json_path='E:\\dataset'        
    xml2json(xml_path,json_path)

    pass


'''
输出效果如下：
loading...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9/9 [00:00<00:00, 143.24it/s]        
9
类名：['red', 'crosswalk']
字典形式：{0: 'red', 1: 'crosswalk'}
'''

classes_indices.json文件如下

在这里插入图片描述

xmlyololabel_121">3、xml转为yolo的label形式

创建xml_to_yolo_label_txt.py,在转换之前，我们需要获取数据集目标类别的列表，即运行上方create_classes_json.py，得到类名集，再替换CLASSES值。程序运行完后自动生成Annotations文件夹，里面就是yolo的label形式txt文件了。

import xml.etree.ElementTree as ET
import os


def convert(size,box):
    # 将bbox的左上角点，右下角点坐标的格式，转换为bbox中心点+bbox的W,H的格式，并进行归一化
    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(xml_path,labels_path,CLASSES,image_id):
    # 把图像image_id的xml文件转换为目标检测的label文件(txt)
    # 其中包含物体的类别cls,bbox的中心点坐标,以及bbox的W,H
    # 并将四个物理量归一化
    in_file=open(xml_path+image_id)
    stats = os.stat(in_file.name)
    if stats.st_size!=0:
        image_id=image_id.split(".")[0]
        out_file=open(labels_path+"%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 代表是否难以识别，0表示易识别，1表示难识别。
            difficult=obj.find("difficult").text
            obj_cls=obj.find("name").text
            
            if obj_cls not in CLASSES:
                continue
            cls_id=CLASSES.index(obj_cls)
            xmlbox=obj.find("bndbox")
            points=(float(xmlbox.find("xmin").text),
                    float(xmlbox.find("xmax").text),
                    float(xmlbox.find("ymin").text),
                    float(xmlbox.find("ymax").text))
            bb=convert((w,h),points)
            out_file.write(str(cls_id)+" "+" ".join([str(a) for a in bb])+"\n")

def make_label_txt(xml_path,labels_path,CLASSES):
    # labels文件夹下创建image_id.txt
    # 对应每个image_id.xml提取出的bbox信息
    image_ids=os.listdir(xml_path)
    for file in image_ids:
        convert_annotation(xml_path,labels_path,CLASSES,file)

if __name__=="__main__":
    # 类别，运行create_classes_json.py，得到类名集
    CLASSES=['red', 'crosswalk']
    # 数据整体文件路径
    common_path='E:\\dataset'
    # xml文件路径
    xml_path=os.path.join(common_path,'xml_outputs\\')
    # Annotations路径，即存放xml转为全部lables的路径
    labels_path=os.path.join(common_path,'Annotations\\')
    if not os.path.exists(labels_path):
        os.mkdir(labels_path)

    # 开始提取和转换
    make_label_txt(xml_path,labels_path,CLASSES)

4、优化代码

有人问，上方两个代码，需要运行两次，能否直接一次性运行就完呢？

答案：是可以的，上方代码关键就在获取类名集，于是将两者合并，创建final_xml2yolo.py代码

from doctest import REPORTING_FLAGS
from lib2to3.pgen2.token import RPAR
import os
from tqdm import tqdm
from lxml import etree
import json
import xml.etree.ElementTree as ET
import os
 
# 读取 xml 文件信息，并返回字典形式
def parse_xml_to_dict(xml):
    if len(xml) == 0:  # 遍历到底层，直接返回 tag对应的信息
        return {xml.tag: xml.text}
 
    result = {}
    for child in xml:
        child_result = parse_xml_to_dict(child)  # 递归遍历标签信息
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:  # 因为object可能有多个，所以需要放入列表里
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}
 
 
# 提取xml中name保留为json文件
def xml2json(data,json_path):
    xml_path = [os.path.join(data, i) for i in os.listdir(data)]
    classes = []      # 目标类别
    num_object = 0
    for xml_file in tqdm(xml_path, desc="loading..."):
        with open(xml_file,encoding='gb18030',errors='ignore') as fid:      # 防止出现非法字符报错
            xml_str = fid.read()
        xml = etree.fromstring(xml_str)
        data = parse_xml_to_dict(xml)["annotation"]  # 读取xml文件信息
        for j in data['object']:        # 获取单个xml文件的目标信息
            ob = j['name']
            num_object +=1
            if ob not in classes:
                classes.append(ob)
    print(num_object)
    # 生成json文件
    labels = {}
    for index,object in enumerate(classes):
        labels[index] = object

    # 打印类名
    classes_name=[labels[key] for key in labels]
    print(f'类名：{classes_name}')
    # 打印类型字典
    print(f'字典形式：{labels}')

    # json.dumps将python对象转为json对象（将dict转化成str）。 json.loads将json字符串解码成python对象（将str转化成dict）
    labels = json.dumps(labels,indent=4)
    json_path=os.path.join(json_path,'classes_indices.json')
    with open(json_path,'w') as f:
        f.write(labels)

    # 返回类名 
    return classes_name


def convert(size,box):
    # 将bbox的左上角点，右下角点坐标的格式，转换为bbox中心点+bbox的W,H的格式，并进行归一化
    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(xml_path,labels_path,CLASSES,image_id):
    # 把图像image_id的xml文件转换为目标检测的label文件(txt)
    # 其中包含物体的类别cls,bbox的中心点坐标,以及bbox的W,H
    # 并将四个物理量归一化
    in_file=open(xml_path+image_id)
    stats = os.stat(in_file.name)
    if stats.st_size!=0:
        image_id=image_id.split(".")[0]
        out_file=open(labels_path+"%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 代表是否难以识别，0表示易识别，1表示难识别。
            difficult=obj.find("difficult").text
            obj_cls=obj.find("name").text
            
            if obj_cls not in CLASSES:
                continue
            cls_id=CLASSES.index(obj_cls)
            xmlbox=obj.find("bndbox")
            points=(float(xmlbox.find("xmin").text),
                    float(xmlbox.find("xmax").text),
                    float(xmlbox.find("ymin").text),
                    float(xmlbox.find("ymax").text))
            bb=convert((w,h),points)
            out_file.write(str(cls_id)+" "+" ".join([str(a) for a in bb])+"\n")

def make_label_txt(xml_path,labels_path,CLASSES):
    # labels文件夹下创建image_id.txt
    # 对应每个image_id.xml提取出的bbox信息
    image_ids=os.listdir(xml_path)
    for file in image_ids:
        convert_annotation(xml_path,labels_path,CLASSES,file)

 
if __name__ == "__main__":
    # 数据整体文件路径，同时也是存放类名json路径
    common_path=json_path='E:\\dataset'
    # xml文件路径
    xml_path=os.path.join(common_path,'xml_outputs\\')
    # 获取类别     
    CLASSES=xml2json(xml_path,json_path)
   
    # Annotations路径，即存放xml转为全部lables的路径
    labels_path=os.path.join(common_path,'Annotations\\')
    if not os.path.exists(labels_path):
        os.mkdir(labels_path)

    # 开始提取和转换
    make_label_txt(xml_path,labels_path,CLASSES)

    pass

5、划分数据集

创建split_train_val.py，根据具体情况分别修改cur_path,image_original_path,label_original_path 值，程序运行完后，分别生成images，labes和data_txt文件夹。

# 将图片和标注数据按比例切分为 训练集和测试集
import shutil
import random
import os

# 原始路径
# 数据整体文件路径(注意这里需要修改)
cur_path='E:\\dataset'
# 图像文件夹路径,注意一定要有\\，(注意这里需要修改)
image_original_path = os.path.join(cur_path,"images_package\\")

# 标注结果的路径即labels路径,该路径下不要有classes.txt,注意一定要有\\，(注意这里需要修改)
label_original_path = os.path.join(cur_path,"Annotations\\")		

# cur_path = os.getcwd()

# 训练集路径
train_image_path = os.path.join(cur_path, "images/train/")
train_label_path = os.path.join(cur_path, "labels/train/")

# 验证集路径
val_image_path = os.path.join(cur_path, "images/val/")
val_label_path = os.path.join(cur_path, "labels/val/")

# 测试集路径
test_image_path = os.path.join(cur_path, "images/test/")
test_label_path = os.path.join(cur_path, "labels/test/")

# 训练集目录
data_txt_path=os.path.join(cur_path,'data_txt')
if not os.path.exists(data_txt_path):
    os.mkdir(data_txt_path)
list_train = os.path.join(data_txt_path, "train.txt")
list_val = os.path.join(data_txt_path, "val.txt")
list_test = os.path.join(data_txt_path, "test.txt")

# 划分数据集比例
train_percent = 0.8
val_percent = 0.1
test_percent = 0.1
 

def del_file(path):
    for i in os.listdir(path):
        file_data = path + "\\" + i
        os.remove(file_data)
 
 
def mkdir():
    if not os.path.exists(train_image_path):
        os.makedirs(train_image_path)
    else:
        del_file(train_image_path)
    if not os.path.exists(train_label_path):
        os.makedirs(train_label_path)
    else:
        del_file(train_label_path)
 
    if not os.path.exists(val_image_path):
        os.makedirs(val_image_path)
    else:
        del_file(val_image_path)
    if not os.path.exists(val_label_path):
        os.makedirs(val_label_path)
    else:
        del_file(val_label_path)
 
    if not os.path.exists(test_image_path):
        os.makedirs(test_image_path)
    else:
        del_file(test_image_path)
    if not os.path.exists(test_label_path):
        os.makedirs(test_label_path)
    else:
        del_file(test_label_path)
 
 
def clearfile():
    if os.path.exists(list_train):
        os.remove(list_train)
    if os.path.exists(list_val):
        os.remove(list_val)
    if os.path.exists(list_test):
        os.remove(list_test)
 
 
def main():
    mkdir()
    clearfile()
 
    file_train = open(list_train, 'w')
    file_val = open(list_val, 'w')
    file_test = open(list_test, 'w')
 
    total_txt = os.listdir(label_original_path)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)
 
    num_train = int(num_txt * train_percent)
    num_val = int(num_txt * val_percent)
    num_test = num_txt - num_train - num_val
 
    train = random.sample(list_all_txt, num_train)
    # train从list_all_txt取出num_train个元素
    # 所以list_all_txt列表只剩下了这些元素
    val_test = [i for i in list_all_txt if not i in train]
    # 再从val_test取出num_val个元素，val_test剩下的元素就是test
    val = random.sample(val_test, num_val)
 
    print("训练集数目：{}, 验证集数目：{}, 测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
    for i in list_all_txt:
        name = total_txt[i][:-4]
 
        srcImage = image_original_path + name + '.jpg'
        srcLabel = label_original_path + name + ".txt"
 
        if i in train:
            dst_train_Image = train_image_path + name + '.jpg'
            dst_train_Label = train_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_train_Image)
            shutil.copyfile(srcLabel, dst_train_Label)
            file_train.write(dst_train_Image + '\n')
        elif i in val:
            dst_val_Image = val_image_path + name + '.jpg'
            dst_val_Label = val_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_val_Image)
            shutil.copyfile(srcLabel, dst_val_Label)
            file_val.write(dst_val_Image + '\n')
        else:
            dst_test_Image = test_image_path + name + '.jpg'
            dst_test_Label = test_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_test_Image)
            shutil.copyfile(srcLabel, dst_test_Label)
            file_test.write(dst_test_Image + '\n')
 
    file_train.close()
    file_val.close()
    file_test.close()
 
if __name__ == "__main__":
    main()

本文最终形成的目录形式

    E:
    └──dataset
        ├── Annotations
        │    ├── 000002.txt
        │    ├── 000003.txt
        │    ├── 000004.txt
        │    ├── 000005.txt
        │    ├── 000006.txt
        │    ├── zebra_crossing_20180129-000545_362.txt
        │    ├── zebra_crossing_20180129-000545_364.txt
        │    ├── zebra_crossing_20180129-000545_366.txt
        │    └── zebra_crossing_20180129-000645_368.txt
        ├── classes_indices.json
        ├── data_txt
        │    ├── test.txt
        │    ├── train.txt
        │    └── val.txt
        ├── images
        │    ├── test
        │    │    ├── 000005.jpg
        │    │    └── zebra_crossing_20180129-000545_362.jpg
        │    ├── train
        │    │    ├── 000002.jpg
        │    │    ├── 000003.jpg
        │    │    ├── 000004.jpg
        │    │    ├── 000006.jpg
        │    │    ├── zebra_crossing_20180129-000545_364.jpg
        │    │    ├── zebra_crossing_20180129-000545_366.jpg
        │    │    └── zebra_crossing_20180129-000645_368.jpg
        │    └── val
        ├── images_package
        │    ├── 000002.jpg
        │    ├── 000003.jpg
        │    ├── 000004.jpg
        │    ├── 000005.jpg
        │    ├── 000006.jpg
        │    ├── zebra_crossing_20180129-000545_362.jpg
        │    ├── zebra_crossing_20180129-000545_364.jpg
        │    ├── zebra_crossing_20180129-000545_366.jpg
        │    └── zebra_crossing_20180129-000645_368.jpg
        ├── labels
        │    ├── test
        │    │    ├── 000005.txt
        │    │    └── zebra_crossing_20180129-000545_362.txt
        │    ├── train
        │    │    ├── 000002.txt
        │    │    ├── 000003.txt
        │    │    ├── 000004.txt
        │    │    ├── 000006.txt
        │    │    ├── zebra_crossing_20180129-000545_364.txt
        │    │    ├── zebra_crossing_20180129-000545_366.txt
        │    │    └── zebra_crossing_20180129-000645_368.txt
        │    └── val
        └── xml_outputs
            ├── 000002.xml
            ├── 000003.xml
            ├── 000004.xml
            ├── 000005.xml
            ├── 000006.xml
            ├── zebra_crossing_20180129-000545_362.xml
            ├── zebra_crossing_20180129-000545_364.xml
            ├── zebra_crossing_20180129-000545_366.xml
            └── zebra_crossing_20180129-000645_368.xml

6、画目录树

创建draw_tree.py，如本文输入

请输入文件夹路径（不含名称）： E

请输入文件夹名称：dataset

自动生成 tree.txt

import os

def get_num(path):
    dirlist = os.listdir(path)
    j=0
    for i in dirlist:
        j+=1
    return j

def print_tree(path,last):
    num=get_num(path)
    if num!=0:
        dirlist = os.listdir(path)
        j=0
        for i in dirlist:
            for k in last:
                if k=='0':
                    print("  │",end=" ")
                else:
                    print("   ", end=" ")
            j+=1
            if j<num:
                print("  ├── ", end="")
                print(i)
                dir=path+"\\"+i
                if os.path.isdir(dir):
                    print_tree(dir,last+'0')
            else:
                print("  └── ", end="")
                print(i)
                dir = path + "\\" + i
                if os.path.isdir(dir):
                    print_tree(dir,last+'1')

def write_tree(path,last,f):
    num=get_num(path)
    if num!=0:
        dirlist = os.listdir(path)
        j=0
        for i in dirlist:
            for k in last:
                if k=='0':
                    f.write("    │")
                else:
                    f.write("    ")
            j+=1
            if j<num:
                f.write("    ├── ")
                f.write(i)
                f.write('\n')
                dir=path+"\\"+i
                if os.path.isdir(dir):
                    write_tree(dir,last+'0',f)
            else:
                f.write("    └── ")
                f.write(i)
                f.write('\n')
                dir = path + "\\" + i
                if os.path.isdir(dir):
                    write_tree(dir,last+'1',f)

if __name__=='__main__':
    path = input("请输入文件夹路径（不含名称）：")
    root = input("请输入文件夹名称：")
    if len(path)==1:
        path+=':'
    #print("  └─root")
    #print_tree('D:\\root',"1")
    f = open("tree.txt", "w", encoding="utf-8")
    f.write("    └──"+root+"\n")
    write_tree(path+"\\"+root, "1",f)
    f.close()