文章目录
1、前言：
本文演示如何划分数据集，以及将VOC标注的xml数据转为YOLO标注的txt格式，且生成classes的txt文件。
python
# 本文演示的项目目录
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.xml2、生成对应的类名
创建
create_classes_json.py自动生成对应的类名json文件，以及在控制台输出对应的类名集。
python
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文件如下
3、xml转为yolo的label形式
创建
xml_to_yolo_label_txt.py,在转换之前，我们需要获取数据集目标类别的列表，即运行上方create_classes_json.py，得到类名集，再替换CLASSES值。程序运行完后自动生成Annotations文件夹，里面就是yolo的label形式txt文件了。
python
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代码
python
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)
pass5、划分数据集
创建
split_train_val.py，根据具体情况分别修改cur_path,image_original_path,label_original_path值，程序运行完后，分别生成images，labes和data_txt文件夹。
python
# 将图片和标注数据按比例切分为 训练集和测试集
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()本文最终形成的目录形式
python
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.xml6、画目录树
创建
draw_tree.py，如本文输入请输入文件夹路径（不含名称）： E
请输入文件夹名称：dataset
自动生成 tree.txt
python
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()