【yolov5】实现一键yolov5数据处理(上)随机抽样+xml转txt+输出量化分析)
·
算法功能:
- 实现从数个数据集目录随机抽取数据
- 可以根据设置classes_name筛选xml里的类,不匹配的类会跳过
- xml转txt(yolo)
- 将所有提取的图像数据存入指定的images目录,将所有提取的txt标签文件存入指定的labels目录
- 输出量化分析表
预先准备:
- 提供一个或者数个已标注过的目录,目录下的图像目录和xml 标签目录一一对应:

然后在代码中输入需要提取的目录地址,提取方式,提取数量,图像格式:
'''
#################### 输入参数设置(开始) ####################
'''
# org_dirs 设置说明:数组内每一个元素表示一条数据集抽离指令, 每一条抽离指令包含4个元素,分别是:
# 目录源路径,计数方式(百分比'rate'/数值'num'),具体的值('rate':[0,1]||'num'>=0), 图像格式
org_dirs = [
['D:\\dataset\\makelabels\\chenlw_221014_hatlungu', 'ratio', '0.8', 'jpg'],
['D:\\dataset\\makelabels\\chenlw_coal', 'num', '500', 'jpg'],
['D:\\dataset\\makelabels\\chenlw_lungu', 'num', '100', 'jpg'],
]
- 设置好其他的提取参数:
# 是否需要统一输出的图像格式
is_union = False
# 若需要统一, 设置统一的图像格式
union_imgtype = 'jpg'
# 输出的目标目录,程序会在该目录下创建images和labels分别存放图像文件和txt文件
output_dir = 'D:\dataset\makelabels'
# 数据集的类设置
classes = ['fire', 'coal', 'hat', 'person', 'Roller']
# 输出量化分析summary.txt(位置在output_dir目录下)
write_summary = True
效果:
会在指定output_dir 目录下生成images和labels目录,分别存储所有提取的图像和标签文件(txt),还会生成summary.txt量化统计文件:


代码:
import os
import numpy as np
import shutil
import xml.etree.ElementTree as ET
import ntpath
'''
算法思想:
1. 将数据从源目录中按比例/数量,随机抽出,最后存放到2个总文件夹中(分别存放xml和jpg).
前提,源目录中必须同时存在jpg文件和xml文件
(若源目录中存在子目录,那么子目录中的文件一样会被提取出来)
2. 在传出过程中将xml文件转为txt文件
更新:
1. labelImg 不需要标注的图像就不会产生xml,因此,对没有匹配到xml的jpg,会自动生成一个空的xml文件,表示无标注
2. 支持判断xml和txt两种标签, xml标签会转txt标签,txt标签会直接复制到指定目录
注意:
凡是图像没有对应的标签,会自动创建空标签!
'''
'''
#################### 输入参数设置(开始) ####################
'''
# org_dirs 设置说明:数组内每一个元素表示一条数据集抽离指令, 每一条抽离指令包含4个元素,分别是:
# 目录源路径,计数方式(百分比'ratio'/数值'num'),具体的值('rate':[0,1]||'num'>=0), 图像格式
org_dirs = [
[r'D:\dataset\02-标注数据集\102-基础数据\BASE\1-fire\BASE-fire-A', 'num', '1000', 'jpg'],
[r'D:\dataset\02-标注数据集\102-基础数据\BASE\1-fire\BASE-fire-B', 'num', '1000', 'jpg'],
[r'D:\dataset\02-标注数据集\102-基础数据\BASE\2-hat\BASE-hat-B-2818', 'num', '2000', 'jpg'],
[r'D:\dataset\02-标注数据集\102-基础数据\BASE\3-person', 'num', '1008', 'jpg'],
[r'D:\dataset\02-标注数据集\101-输煤栈桥\1-roller\2-大连奥通\JG-roller\DLAT-hat-roller-v1-594', 'num', '594', 'jpg'],
[r'D:\dataset\02-标注数据集\101-输煤栈桥\1-roller\2-大连奥通\JG-roller\DLAT-roller-v1-968', 'num', '968', 'jpg'],
]
# 是否需要统一输出的图像格式
is_union = False
# 若需要统一, 设置统一的图像格式
union_imgtype = 'jpg'
# 输出的目标目录,程序会在该目录下创建images和labels分别存放图像文件和txt文件
output_dir = r'D:\dataset\test'
#巨哥相机数据集规则
classes = ['fire', 'Roller', 'hat', 'person']
# 输出量化分析summary.txt(位置在output_dir目录下)
write_summary = True
'''
#################### 输入参数设置(结束) ####################
'''
# txt数据统计设置
summary_dataset=[]
summary_cls_to_img = []
summary_cls_to_obj = []
summary_img = 0
summary_obj = 0
for _ in classes:
summary_cls_to_img.append(0)
summary_cls_to_obj.append(0)
# 按照yolo格式归一化
def convert(size, box):
dw = 1.0 / size[0]
dh = 1.0 / 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 count_txt(txt_path):
flag_cls = []
for _ in classes:
flag_cls.append(0)
with open(txt_path, 'r',encoding='UTF-8') as f:
while True:
line = f.readline()
if line:
# 以空格为间隔,读取所有该行的数据存入数组
msg = line.split(" ")
if len(msg) == 5:
cls_idx = line.split(' ')[0]
# print("cls_idx=", cls_idx)
summary_cls_to_obj[int(cls_idx)] += 1
flag_cls[int(cls_idx)] = 1
else:
break
for idx in range(len(classes)):
if flag_cls[idx] == 1:
summary_cls_to_img[idx] += 1
# xml转txt(yolo)
def convert_annotation(xml_path, classes, summary_obj):
flag_cls = []
for _ in classes:
flag_cls.append(0)
# 读取xml文件,写入txt文件
# print(xml_path)
# 判断xml是否为空
if not os.path.getsize(xml_path):
print("内容为空的xml:",xml_path)
txt_path = xml_path[:-3] + 'txt' # xml文件路径
with open(txt_path, 'w') as out_txt:
return summary_obj, txt_path
with open(xml_path,encoding='UTF-8') as f:
xml_text = f.read()
# tree = ET.parse(f)
# root = tree.getroot()
root = ET.fromstring(xml_text)
f.close()
# 获取xml参数信息
outlines = []
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
cls = obj.find('name').text
if cls not in classes:
print('invalid_cls:',cls)
continue
summary_obj += 1
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)
summary_cls_to_obj[cls_id] += 1
flag_cls[cls_id] = 1
outlines.append(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
# summary
for idx in range(len(classes)):
if flag_cls[idx] == 1:
summary_cls_to_img[idx] += 1
# 写入txt文件
txt_path = xml_path[:-3] + 'txt' # xml文件路径
with open(txt_path, 'w') as out_txt:
out_txt.writelines(outlines)
return summary_obj, txt_path
if __name__ == "__main__":
# copy file and convert xml to yolo
# 创建目录
if not os.path.exists(output_dir):
os.mkdir(output_dir)
output_images = os.path.join(output_dir, 'images')
if not os.path.exists(output_images):
os.mkdir(output_images)
output_labels = os.path.join(output_dir, 'labels')
if not os.path.exists(output_labels):
os.mkdir(output_labels)
# 遍历目录
for dir in org_dirs:
# create dir
path, mode, val, images_type = dir
old_path = []
for root, dir, files in os.walk(path):
for file in files:
if file[-3:] == images_type:
# print(file)
xml_name = file[:-3] + 'xml'
txt_name = file[:-3] + 'txt'
if not (os.path.exists(os.path.join(root, xml_name)) or os.path.exists(os.path.join(root, txt_name))):
# 生成空的txt文件
f = open(os.path.join(root, txt_name), 'w', encoding='utf-8')
f.close()
old_path.append(os.path.join(root, file))
if len(old_path) <= 0:
print('目录:', path, '为空,跳过!\n')
continue
# 复制到 output_dir
if mode == 'ratio': # 按比例抽取
if float(val) >=0 and float(val) <= 1:
num = int(len(old_path) * float(val))
else:
print('提示:目录:', path, 'val=', val,',无效,跳过!\n')
continue
elif mode == 'num': # 按数值抽取
if int(val) > len(old_path):
print('提示:目录:', path, 'val=', val,',超过目录内目标上线,默认提取目录内全部有效数据 num=',len(old_path),'\n')
num = len(old_path)
elif int(val) <= len(old_path) and int(val) > 0:
num = int(val)
else:
print('提示:目录:', path, 'val=', val, ',无效,跳过!\n')
continue
# 打乱序列
np.random.shuffle(old_path)
count=0
for p in old_path[:num]:
count+= 1
summary_img += 1
# 图像复制
imgname = ntpath.basename(p)
shutil.copyfile(p, os.path.join(output_images,imgname))
xmlpath = p[:-3] + 'xml'
if os.path.exists(xmlpath):
summary_obj, txt_path =convert_annotation(xmlpath, classes, summary_obj)
# txt移动
txt_name = imgname[:-3] + 'txt'
shutil.move(txt_path, os.path.join(output_labels, txt_name))
# print('目录',path,'转换复制进度:',count,'/',str(num))
elif os.path.exists(xmlpath[:-3]+"txt"):
txt_path = xmlpath[:-3] + 'txt'
txt_name = imgname[:-3] + 'txt'
shutil.copyfile(txt_path, os.path.join(output_labels, txt_name))
count_txt(txt_path)
dataset_status = '数据集名称:' + path + ',随机抽样(' + str(num) + '/' + str(len(old_path)) + ')\n'
summary_dataset.append(dataset_status)
print('目录',path,'完成数据提取!' )
# 量化分析
if write_summary:
txt_path = os.path.join(output_dir,'summary.txt')
with open(txt_path, 'w') as summary_txt:
summary_txt.write('数据集抽样统计:\n')
for line in summary_dataset:
summary_txt.write(line)
# 编写类别编号
classes_line = 'classes=['
for idx in range(len(classes)):
cls_msg = '\''+classes[idx] + '\''
classes_line += cls_msg
if idx < len(classes)-1:
classes_line += ','
classes_line += ']\n'
summary_txt.write(classes_line)
summary_txt.write('\n量化分析:\n')
summary_txt.write('类在图像中出现的频率统计:\n')
for idx in range(len(classes)):
ratio = float(summary_cls_to_img[idx] / summary_img)
mes = classes[idx] + ':出现次数=' + str(summary_cls_to_img[idx]) + ',比例=' + str(round(ratio, 4)) + '\n'
summary_txt.write(mes)
summary_txt.write('\n类在检测框中出现的频率统计:\n')
for idx in range(len(classes)):
ratio = float(summary_cls_to_obj[idx] / summary_obj)
mes = classes[idx] + ':出现次数=' + str(summary_cls_to_obj[idx]) + ',比例=' + str(round(ratio, 4)) + '\n'
summary_txt.write(mes)
print( '完成量化分析!')
更多推荐


所有评论(0)