数据集扩充

前言

为增加数据集中图片的数量通常需要将几个数据集整合在一起，但每个数据集的格式或数据集中标签可能与自己的不尽相同，也是搜寻了好几天才好到几个好用的，怕以后忘了，所以记录一下。

---

找了好几个博主的代码试了一下，都搞混了。但是还是在此感谢一下CSDN上的各位大佬！！

一、yolo格式的标签txt转VOC格式的标签xml？

我的txt文件格式如下，如果你的不是我也不知道咋办！
1 0.8125 0.45583333333333337 0.08 0.09833333333333334
0 0.9175 0.47250000000000003 0.0675 0.08833333333333333
只需要改一下字典和文件路径（路径中不要有中文）

代码如下：

from xml.dom.minidom import Document
import os
import cv2


def makexml(txtPath, xmlPath, picPath):  # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径
    """此函数用于将yolo格式txt标注文件转换为voc格式xml标注文件
    """
    dic = {'0': "nomsk",  # 创建字典用来对类型进行转换
           '1': "mask"  # 此处的字典要与自己的classes.txt文件中的类对应，且顺序要一致
           }
    files = os.listdir(txtPath)
    for i, name in enumerate(files):
        xmlBuilder = Document()
        annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签
        xmlBuilder.appendChild(annotation)
        txtFile = open(txtPath + name)
        txtList = txtFile.readlines()
        img = cv2.imread(picPath + name[0:-4] + ".jpg")
        Pheight, Pwidth, Pdepth = img.shape

        folder = xmlBuilder.createElement("folder")  # folder标签
        foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset")
        folder.appendChild(foldercontent)
        annotation.appendChild(folder)  # folder标签结束

        filename = xmlBuilder.createElement("filename")  # filename标签
        filenamecontent = xmlBuilder.createTextNode(name[0:-4] + ".jpg")
        filename.appendChild(filenamecontent)
        annotation.appendChild(filename)  # filename标签结束

        size = xmlBuilder.createElement("size")  # size标签
        width = xmlBuilder.createElement("width")  # size子标签width
        widthcontent = xmlBuilder.createTextNode(str(Pwidth))
        width.appendChild(widthcontent)
        size.appendChild(width)  # size子标签width结束

        height = xmlBuilder.createElement("height")  # size子标签height
        heightcontent = xmlBuilder.createTextNode(str(Pheight))
        height.appendChild(heightcontent)
        size.appendChild(height)  # size子标签height结束

        depth = xmlBuilder.createElement("depth")  # size子标签depth
        depthcontent = xmlBuilder.createTextNode(str(Pdepth))
        depth.appendChild(depthcontent)
        size.appendChild(depth)  # size子标签depth结束

        annotation.appendChild(size)  # size标签结束

        for j in txtList:
            oneline = j.strip().split(" ")
            object = xmlBuilder.createElement("object")  # object 标签
            picname = xmlBuilder.createElement("name")  # name标签
            namecontent = xmlBuilder.createTextNode(dic[oneline[0]])
            picname.appendChild(namecontent)
            object.appendChild(picname)  # name标签结束

            pose = xmlBuilder.createElement("pose")  # pose标签
            posecontent = xmlBuilder.createTextNode("Unspecified")
            pose.appendChild(posecontent)
            object.appendChild(pose)  # pose标签结束

            truncated = xmlBuilder.createElement("truncated")  # truncated标签
            truncatedContent = xmlBuilder.createTextNode("0")
            truncated.appendChild(truncatedContent)
            object.appendChild(truncated)  # truncated标签结束

            difficult = xmlBuilder.createElement("difficult")  # difficult标签
            difficultcontent = xmlBuilder.createTextNode("0")
            difficult.appendChild(difficultcontent)
            object.appendChild(difficult)  # difficult标签结束

            bndbox = xmlBuilder.createElement("bndbox")  # bndbox标签
            xmin = xmlBuilder.createElement("xmin")  # xmin标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)
            xminContent = xmlBuilder.createTextNode(str(mathData))
            xmin.appendChild(xminContent)
            bndbox.appendChild(xmin)  # xmin标签结束

            ymin = xmlBuilder.createElement("ymin")  # ymin标签
            mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)
            yminContent = xmlBuilder.createTextNode(str(mathData))
            ymin.appendChild(yminContent)
            bndbox.appendChild(ymin)  # ymin标签结束

            xmax = xmlBuilder.createElement("xmax")  # xmax标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)
            xmaxContent = xmlBuilder.createTextNode(str(mathData))
            xmax.appendChild(xmaxContent)
            bndbox.appendChild(xmax)  # xmax标签结束

            ymax = xmlBuilder.createElement("ymax")  # ymax标签
            mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)
            ymaxContent = xmlBuilder.createTextNode(str(mathData))
            ymax.appendChild(ymaxContent)
            bndbox.appendChild(ymax)  # ymax标签结束

            object.appendChild(bndbox)  # bndbox标签结束

            annotation.appendChild(object)  # object标签结束

        f = open(xmlPath + name[0:-4] + ".xml", 'w')
        xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
        f.close()


makexml("D:/DeepLearning/data/txt/", "D:/DeepLearning/data/xml/",
        "D:/DeepLearning/data/img/")  # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径,注意不要有中文

二、更改数据集中的标签

1.代码如下（示例）：

import os
import xml.etree.ElementTree as ET

#程序功能：批量修改VOC数据集中xml标签文件的标签名称
def changelabelname(inputpath):
    listdir = os.listdir(inputpath)
    for file in listdir:
        if file.endswith('xml'):
            file = os.path.join(inputpath,file)
            tree = ET.parse(file)
            root = tree.getroot()
            for object1 in root.findall('object'):
                for sku in object1.findall('name'):           #查找需要修改的名称
                    if (sku.text == '1'):               #‘preName’为修改前的名称
                        sku.text = 'nomask'                 #‘TESTNAME’为修改后的名称
                        tree.write(file,encoding='utf-8')     #写进原始的xml文件并避免原始xml中文字符乱码
                    else:
                        pass
        else:
            pass

if __name__ == '__main__':
    inputpath = 'D:/DeepLearning/VOCdevkit（训练集）/VOCdevkit/VOC2007/Annotations'  #此处替换为自己的路径
    changelabelname(inputpath)

缺点是每次只能更改1个标签

统计VOC数据集中的类别个数

#VOC 数据集类别统计
import xml.dom.minidom as xmldom
import os


def voc_label_statistics(annotation_path):
    '''
    VOC 数据集类别统计
    :param annotation_path: voc数据集的标签文件夹
    :return: {'class1':'count',...}
    '''
    count = 0
    annotation_names = [os.path.join(annotation_path, i) for i in os.listdir(annotation_path)]

    labels = dict()
    for names in annotation_names:
        names_arr = names.split('.')
        file_type = names_arr[-1]
        if file_type != 'xml':
            continue
        file_size = os.path.getsize(names)
        if file_size == 0:
            continue

        count = count + 1
        print('process：', 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

            label_count = labels.get(label, 0)
            labels[label] = label_count + 1

    print('文件标注个数：', count)
    return labels


if __name__ == '__main__':
    annotation_path = "D:/BaiduNetdiskDownload/烟雾检测数据集/xml"

    label = voc_label_statistics(annotation_path)
    print(label)

批量移除xml标注中的某一个类别标签

#  批量移除xml标注中的某一个类别标签
import xml.etree.cElementTree as ET
import os

# xml文件路径
xml_path = r'D:/BaiduNetdiskDownload/烟雾检测数据集/xml'
xml_files = os.listdir(xml_path)
# 需要删除的类别名称
CLASSES = ["fire"]

for axml in xml_files:
    path_xml = os.path.join(xml_path, axml)
    tree = ET.parse(path_xml)
    root = tree.getroot()
    for child in root.findall('object'):
        name = child.find('name').text
        if name in CLASSES:
            root.remove(child)
    tree.write(os.path.join(xml_path, axml))

根据VOC数据集中的xml文件在图像上绘制ground-boxes

import os
import os.path
import numpy as np
import xml.etree.ElementTree as xmlET
from PIL import Image, ImageDraw

classes = ('__background__', # always index 0
           'smoke')

file_path_img = 'D:/BaiduNetdiskDownload/day_time_wildfire_v2/images'
file_path_xml = 'D:/BaiduNetdiskDownload/day_time_wildfire_v2/annotations/xmls'
save_file_path = 'D:/BaiduNetdiskDownload/day_time_wildfire_v2/groundtruth'

pathDir = os.listdir(file_path_xml)
for idx in range(len(pathDir)):  
    filename = pathDir[idx]
    tree = xmlET.parse(os.path.join(file_path_xml, filename))
    objs = tree.findall('object')        
    num_objs = len(objs)
    boxes = np.zeros((num_objs, 5), dtype=np.uint16)

    for ix, obj in enumerate(objs):
        bbox = obj.find('bndbox')
        # Make pixel indexes 0-based
        x1 = float(bbox.find('xmin').text) - 1
        y1 = float(bbox.find('ymin').text) - 1
        x2 = float(bbox.find('xmax').text) - 1
        y2 = float(bbox.find('ymax').text) - 1

        cla = obj.find('name').text 
        label = classes.index(cla)

        boxes[ix, 0:4] = [x1, y1, x2, y2]
        boxes[ix, 4] = label

    image_name = os.path.splitext(filename)[0]
    img = Image.open(os.path.join(file_path_img, image_name + '.jpeg')) 

    draw = ImageDraw.Draw(img)
    for ix in range(len(boxes)):
        xmin = int(boxes[ix, 0])
        ymin = int(boxes[ix, 1])
        xmax = int(boxes[ix, 2])
        ymax = int(boxes[ix, 3])
        draw.rectangle([xmin, ymin, xmax, ymax], outline=(255, 0, 0))
        draw.text([xmin, ymin], classes[boxes[ix, 4]], (255, 0, 0))

    img.save(os.path.join(save_file_path, image_name + '.jpeg'))

更改xml文件中的path或filename

# 更改.xml的filename的代码块
import os
import xml.dom.minidom

path='D:/BaiduNetdiskDownload/烟雾检测数据集/VOCdevkit/VOC2007/Annotations' # xml文件存放路径
sv_path='D:/BaiduNetdiskDownload/烟雾检测数据集/VOCdevkit/VOC2007/xml' # 修改后的xml文件存放路径
files=os.listdir(path)

for xmlFile in files:
    dom=xml.dom.minidom.parse(os.path.join(path,xmlFile)) #打开xml文件，送到dom解析
    root=dom.documentElement #得到文档元素对象
    item=root.getElementsByTagName('filename') #获取filename这一node名字及相关属性值  # 想更改path 把filename替换成path 同时把下面的i.firstChild.data的输入替换成路径
    a,b=os.path.splitext(xmlFile) #分离出文件名a
    for i in item:
        i.firstChild.data = a + '.jpg'
    with open(os.path.join(sv_path,xmlFile),'w') as fh:
        dom.writexml(fh)