算法功能:

  1. 实现从数个数据集目录随机抽取数据
  2. 可以根据设置classes_name筛选xml里的类,不匹配的类会跳过
  3. xml转txt(yolo)
  4. 将所有提取的图像数据存入指定的images目录,将所有提取的txt标签文件存入指定的labels目录
  5. 输出量化分析表

预先准备:

  1. 提供一个或者数个已标注过的目录,目录下的图像目录和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'],
]
  1. 设置好其他的提取参数:
# 是否需要统一输出的图像格式
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( '完成量化分析!')

更多推荐