【代码】处理数据集的一些常用脚本

前言
由于训练模型经常使用到一些处理数据集的脚本，每次重新写都蛮麻烦的，因此决定整理一下之后备用，加快效率。一些是借鉴网上的脚本但是因为时间太久找不到原链接所以比较抱歉没有贴出来原链接，还有一些是自己写的，发出来分享一下。

1.voc数据集抽取特定类

import os
import shutil

ann_filepath = 'E:/Construction_Site_of_Intelligentization/Reflective_vests-ori/VOC2021/VOC2021/voc_kk_rename/Annotations/'
img_filepath = 'E:/Construction_Site_of_Intelligentization/Reflective_vests-ori/VOC2021/VOC2021/voc_kk_rename/JPEGImages/'
img_savepath = 'E:/Construction_Site_of_Intelligentization/Reflective_vests-ori/VOC2021/VOC2021/chouqu1/JPEGImages/'
ann_savepath = 'E:/Construction_Site_of_Intelligentization/Reflective_vests-ori/VOC2021/VOC2021/chouqu1/Annotations/'
if not os.path.exists(img_savepath):
    os.mkdir(img_savepath)

if not os.path.exists(ann_savepath):
    os.mkdir(ann_savepath)
names = locals()
classes = ['hat', 'other_clothes', 'head', 'reflective_clothes']  # 总类别

for file in os.listdir(ann_filepath):
    print(file)

    fp = open(ann_filepath + '\\' + file)  # 打开Annotations文件
    ann_savefile = ann_savepath + file
    fp_w = open(ann_savefile, 'w')
    lines = fp.readlines()

    ind_start = []
    ind_end = []
    lines_id_start = lines[:]

    lines_id_end = lines[:]
    
    # 其中要抽取的两类other_clothes和reflective_clothes
    classes1 = '\t\t<name>other_clothes</name>\n'  
    classes2 = '\t\t<name>reflective_clothes</name>\n'
    # classes3 = '\t\t<name>car</name>\n'
    # classes4 = '\t\t<name>motorbike</name>\n'
    # classes5 = '\t\t<name>train</name>\n'

    # 在xml中找到object块，并将其记录下来
    while "\t<object>\n" in lines_id_start:
        a = lines_id_start.index("\t<object>\n")
        ind_start.append(a)  # ind_start是<object>的行数
        lines_id_start[a] = "delete"

    while "\t</object>\n" in lines_id_end:
        b = lines_id_end.index("\t</object>\n")
        ind_end.append(b)  # ind_end是</object>的行数
        lines_id_end[b] = "delete"

    # names中存放所有的object块
    i = 0
    for k in range(0, len(ind_start)):
        names['block%d' % k] = []
        for j in range(0, len(classes)):
            if classes[j] in lines[ind_start[i] + 1]:
                a = ind_start[i]
                for o in range(ind_end[i] - ind_start[i] + 1):
                    names['block%d' % k].append(lines[a + o])
                break
        i += 1
        # print(names['block%d' % k])

    # xml头
    string_start = lines[0:ind_start[0]]

    # xml尾
    if ((file[2:4] == '09') | (file[2:4] == '10') | (file[2:4] == '11')):
        string_end = lines[(len(lines) - 11):(len(lines))]
    else:
        string_end = [lines[len(lines) - 1]]

        # 在给定的类中搜索，若存在则，写入object块信息
    a = 0
    for k in range(0, len(ind_start)):
        if classes1 in names['block%d' % k]:
            a += 1
            string_start += names['block%d' % k]
        if classes2 in names['block%d' % k]:
            a += 1
            string_start += names['block%d' % k]
        # if classes3 in names['block%d' % k]:
        #     a += 1
        #     string_start += names['block%d' % k]
        # if classes4 in names['block%d' % k]:
        #     a += 1
        #     string_start += names['block%d' % k]
        # if classes5 in names['block%d' % k]:
        #     a += 1
        #     string_start += names['block%d' % k]

    string_start += string_end
    # print(string_start)
    for c in range(0, len(string_start)):
        fp_w.write(string_start[c])
    fp_w.close()
    # 如果没有我们寻找的模块，则删除此xml，有的话拷贝图片
    if a == 0:
        os.remove(ann_savepath + file)
    else:
        name_img = img_filepath + os.path.splitext(file)[0] + ".jpg"
        shutil.copy(name_img, img_savepath)
    fp.close()

2.txt格式数据集抽取特定类

import shutil
import os
# classes = ['hat', 'other_clothes', 'head', 'reflective_clothes'] 抽1和3

#输入原始图片和标签的地址
img_path = 'E:/Construction_Site_of_Intelligentization/Reflective_vests/images/'
label_path = 'E:/Construction_Site_of_Intelligentization/Reflective_vests/labels/'
#输入目的图片和标签的地址
new_img_path = 'E:/Construction_Site_of_Intelligentization/Reflective_vests/labels_cq/'
new_label_path = 'E:/Construction_Site_of_Intelligentization/Reflective_vests/images_cq/'

for file in os.listdir(label_path):
    print(file)
    if file == "classes.txt":
        continue
    print(label_path + file)
    with open(label_path + file, 'r') as fp:
        for line in fp.readlines():
            if line.split(' ')[0] == '1' or line.split(' ')[0] == '3':
                shutil.copy(label_path + file, new_label_path)
                img_file = file.split('.')[0] + '.jpg'
                shutil.copy(img_path  + img_file, new_img_path)

3.根据文件名抽取图片

import os
import shutil

file_path = r'E:\Construction Site of Intelligentization\helmet\hat\test.txt'
Targetfile_path = 'E:/Construction Site of Intelligentization/helmet/hat/test/'
with open(file_path, 'r') as f:
    for line in f.readlines():
        name = line.split('/')[-1].strip('\n')
        name = 'E:/Construction Site of Intelligentization/helmet/hat/images/' + name
        shutil.copy(name, Targetfile_path)

4.合并两个文件夹中同名的txt标签文件（用于将同一张图像的不同类别的标签汇总）

# 把kk2文件夹中写到kk1文件夹中同名的

import os

path_ori = 'E:/Construction_Site_of_Intelligentization/kksmoke/labels/'

for filename in os.listdir(path_ori):
    filepath = path_ori + filename
    kkpath = 'E:/Construction_Site_of_Intelligentization/kksmoke/labels-fire/' + filename

    listkk = []
    try:
        with open(kkpath, 'r') as fk:
            for line in fk.readlines():
                listkk.append(line)


        with open(filepath, 'a') as f:
            f.write('\n')
            for i in listkk:
                f.write(i)
    except:
        continue

5.重命名/删除voc数据集的类别

import os
import xml.etree.ElementTree as ET

origin_ann_dir = 'E:/Construction Site of Intelligentization/Reflective_vests-ori/VOC2021/VOC2021/Annotations-ori/'  # 设置原始标签路径为 Annos
new_ann_dir = 'E:/Construction Site of Intelligentization/Reflective_vests-ori/VOC2021/VOC2021/Annotations/'  # 设置新标签路径 Annotations
for dirpaths, dirnames, filenames in os.walk(origin_ann_dir):  # os.walk游走遍历目录名
    for filename in filenames:
        print("process...")
        if os.path.isfile(r'%s%s' % (origin_ann_dir, filename)):  # 获取原始xml文件绝对路径，isfile()检测是否为文件 isdir检测是否为目录
            origin_ann_path = os.path.join(r'%s%s' % (origin_ann_dir, filename))  # 如果是，获取绝对路径（重复代码）
            new_ann_path = os.path.join(r'%s%s' % (new_ann_dir, filename))
            tree = ET.parse(origin_ann_path)  # ET是一个xml文件解析库，ET.parse（）打开xml文件。parse--"解析"
            root = tree.getroot()  # 获取根节点
            for object in root.findall('object'):  # 找到根节点下所有“object”节点
                name = str(object.find('name').text)  # 找到object节点下name子节点的值（字符串）
                # 如果name等于str，则删除该节点
                # if (name in ["person"]):
                #     root.remove(object)

                # # 如果name等于str，则修改name
                if (name in ["person"]):
                    object.find('name').text = "head"

            # # 检查是否存在labelmap中没有的类别
            # for object in root.findall('object'):
            #     name = str(object.find('name').text)
            #     if not (name in ["chepai", "chedeng", "chebiao"]):
            #         print(filename + "------------->label is error--->" + name)
            tree.write(new_ann_path)  # tree为文件，write写入新的文件中。

6.批量修改标签和图片名

import shutil
import os

#输入原始图片和标签的地址
img_path = 'E:/Construction_Site_of_Intelligentization/fire_and_smoke_kk/fire_images/'
label_path = 'E:/Construction_Site_of_Intelligentization/fire_and_smoke_kk/fire_annotations/'
#输入目的图片和标签的地址
new_img_path = 'E:/Construction_Site_of_Intelligentization/fire_and_smoke_kk/new_fire_images/'
new_label_path = 'E:/Construction_Site_of_Intelligentization/fire_and_smoke_kk/new_fire_annotations/'
if not os.path.exists(new_img_path):
    os.mkdir(new_img_path)
if not os.path.exists(new_label_path):
    os.mkdir(new_label_path)


#保存原始的图片和标签名字（带后缀名）
img_list = os.listdir(img_path)
label_list = os.listdir(label_path)
#保存原始的图片和标签名字（不带后缀名）
img_name = []
label_name = []

#获得不带后缀名的图片名字
for img in img_list:
    a, b = img.split('.')
    img_name.append(a)
#获得不带后缀名的标签名字
for label in label_list:
    a, b = label.split('.')
    label_name.append(a)

i = 0
for img in img_name:
    for label in label_name:
        #找到对应的图片和标签
        if img == label:
            n = 6 - len(str(i))
            #原始地址
            img_src = os.path.join(img_path, img+ '.jpg')
            label_src = os.path.join(label_path, label+ '.xml') # txt格式就把这里的xml改成txt
            #目的地址
            # 这里由于类别过多因此加了前缀“fire_”，自用可去掉
            img_dst = os.path.join(new_img_path, 'fire_'+n*str(0)+str(i)+".jpg") 
            label_dst = os.path.join(new_label_path, 'fire_'+n*str(0)+str(i)+".xml")
            #改名并保存倒新的地址
            shutil.copy(img_src, img_dst)
            shutil.copy(label_src, label_dst)
            i = i + 1

7.重命名txt格式标签的类别，即第一列的数字

import os

fpath = r"E:\Construction_Site_of_Intelligentization\kksmoke\labels-ori\labels"  # 读取文件夹路径


def filechanger(path):
    # 读取文件夹中的文件（包括子文件夹、子文件）
    filenames = os.listdir(path)
    for filename in filenames:
        domain = os.path.abspath(path)
        filename = os.path.join(domain, filename)
        if os.path.isdir(filename):
            filechanger(filename)
            continue
        fread = open(filename, 'r')
        fwrite = open("%s.backup" % filename, 'w')
        while True:
            line = fread.readline()
            if len(line) > 0:
                content = line.split()
                # 按照需要修改下面代码
                if content[0] == '0':
                    content[0] = str(4)
                # elif content[0] == '1':
                #     content[0] = str(3)
                # elif content[0] == '2':
                #     content[0] = str(1)
                # elif content[0] == '3':
                #     content[0] = str(0)
                else:
                    content[0] = content[0]
                newcont = content[0] + ' ' + content[1] + ' ' + content[2] + ' ' + content[3] + ' ' + content[4] + '\r'
                fwrite.write(newcont)  # 修改后写入新文件
            else:
                break
        fread.close()
        fwrite.close()
        os.remove(filename)
        os.rename("%s.backup" % filename, filename)


filechanger(fpath)

8.按比例划分数据集（这里演示8：1：1）

import os
import random
# import argparse

# parser = argparse.ArgumentParser()
# # xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
# parser.add_argument('--xml_path', default='Annotations/', type=str, help='input xml label path')
# # 数据集的划分，地址选择自己数据下的ImageSets/Main
# parser.add_argument('--txt_path', default='Main/', type=str, help='output txt label path')
# opt = parser.parse_args()

train_percent = 0.8  # 训练集所占比例
val_percent = 0.1  # 验证集所占比例
test_persent = 0.1  # 测试集所占比例

filepath = r'E:\Construction_Site_of_Intelligentization\kksmoke\images'
txtsavepath = r'E:\Construction_Site_of_Intelligentization\kksmoke\labels'
total_labels = os.listdir(filepath)

if not os.path.exists(txtsavepath):
    os.makedirs(savepath)

num = len(total_labels)
list = list(range(num))

t_train = int(num * train_percent)
t_val = int(num * val_percent)

train = random.sample(list, t_train)
num1 = len(train)
for i in range(num1):
    list.remove(train[i])

val_test = [i for i in list if not i in train]
val = random.sample(val_test, t_val)
num2 = len(val)
for i in range(num2):
    list.remove(val[i])

file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')

for i in train:
    name = total_labels[i] + '\n'
    file_train.write(name)

for i in val:
    name = total_labels[i] + '\n'
    file_val.write(name)

for i in list:
    name = total_labels[i] + '\n'
    file_test.write(name)

file_train.close()
file_val.close()
file_test.close()

9.获得voc数据集所有类别

import xml.dom.minidom as xmldom
import os

annotation_path = r"E:/Construction_Site_of_Intelligentization/fire_and_smoke_kk/smoke_annotations/"

annotation_names=[os.path.join(annotation_path,i) for i in os.listdir(annotation_path)]

labels = list()
for names in annotation_names:
    xmlfilepath = names
    domobj = xmldom.parse(xmlfilepath)
    # 得到元素对象
    elementobj = domobj.documentElement
    #获得子标签
    subElementObj = elementobj.getElementsByTagName("object")
    for s in subElementObj:
        label=s.getElementsByTagName("name")[0].firstChild.data
        #print(label)
        if label not in labels:
            labels.append(label)
print(labels)

10.voc格式转txt格式标签

import os.path
import xml.etree.ElementTree as ET

class_names = ['smoke']

xmlpath = 'E:/Construction_Site_of_Intelligentization/fire_and_smoke_kk/new_smoke_annotations/'  # 原xml路径
txtpath = 'E:/Construction_Site_of_Intelligentization/fire_and_smoke_kk/smoke_labels/'  # 转换后txt文件存放路径
files = []

for root, dirs, files in os.walk(xmlpath):
    None

number = len(files)
print(number)
i = 0
while i < number:

    name = files[i][0:-4]
    xml_name = name + ".xml"
    txt_name = name + ".txt"
    xml_file_name = xmlpath + xml_name
    txt_file_name = txtpath + txt_name

    xml_file = open(xml_file_name, encoding='utf-8')
    tree = ET.parse(xml_file)
    root = tree.getroot()
    filename = root.find('filename').text

    image_name = root.find('filename').text
    w = int(root.find('size').find('width').text)
    h = int(root.find('size').find('height').text)

    f_txt = open(txt_file_name, 'w+')
    content = ""

    first = True

    for obj in root.iter('object'):

        name = obj.find('name').text
        class_num = class_names.index(name)

        xmlbox = obj.find('bndbox')

        x1 = int(xmlbox.find('xmin').text)
        x2 = int(xmlbox.find('xmax').text)
        y1 = int(xmlbox.find('ymin').text)
        y2 = int(xmlbox.find('ymax').text)

        if first:
            content += str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)
            first = False
        else:
            content += "\n" + \
                       str(class_num) + " " + \
                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \
                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)

    # print(str(i / (number - 1) * 100) + "%\n")

    print(content)
    f_txt.write(content)
    f_txt.close()
    xml_file.close()
    i += 1


print("done!")