自动化迁移和更新物体检测XML数据集

引言

当处理大规模的图片数据集，尤其是在物体识别和计算机视觉领域时，有效管理和更新数据集变得非常重要。在本文中，我们将介绍一个简单的Python脚本，该脚本自动化了迁移和更新标注过的XML数据集的过程。

脚本概览

该脚本进行了以下操作：

遍历给定目录下的所有XML注释文件。
为每个注释文件创建对应的拷贝，并更新文件路径和名字。
单独处理每个XML文件中的物体注释，保持标注信息不变。
拷贝原图像文件到新的目录，并更新文件名。
在终端输出处理进度，包括拷贝的图片和XML文件计数器。
把独特的标签类别以及总计数保存到文本文件中。

使用场景

更新数据集目录结构。
从图片中提取注释信息，用于机器学习模型的训练。
快速获取数据集的整体信息，包括包含了哪些类别。

代码解析

设置

在脚本的开始部分，我们首先导入必要的模块，并定义源目录和目标目录的路径。

复制代码

import xml.etree.ElementTree as ET
from datetime import datetime
import os, cv2, time, shutil

这些模块包括操作XML文件的ElementTree，进行文件操作的os和shutil，还有处理图片的cv2等。

功能函数

脚本中定义了两个重要的函数来处理XML文件的创建和物体标记的添加。

复制代码

def create_object(root, xi, yi, xa, ya, obj_name):
    # 创建物体标记代码
    # ...

def create_tree(sources, image_name, h, w):
    # 创建XML文件树结构
    # ...

主循环

在主循环中，脚本读取每个注释文件，并对每个文件执行一系列操作：

复制代码

for xml in annotation_list:
    # ... 每个文件的处理逻辑...

这包括通过OpenCV获取图片尺寸，使用shutil.copyfile拷贝图片文件，解析XML文件以读取和复制物体标注信息。

结果记录

脚本最后输出处理的总数，并将数据类别及其计数保存在文本文件中：

复制代码

print(count)
labels_list = list(set(all_labels))
labels_len = len(labels_list)
labels_str = f"this datasets have {labels_len} classes, labels list is {labels_list}"
print(labels_str)
with open("record.txt", "w") as f:
    f.write(str(count))
    f.write("\n")
    f.write(labels_str)

结论

我们介绍的脚本可以帮助用户以一种简单高效的方式来处理和准备用于物体检测的图片数据集。通过使用此脚本，用户可以节省宝贵的时间并减少手动错误，提高了数据准备工作的效率。

代码

完整脚本可以在本文附带的代码段中找到。如果你有兴趣自己尝试或者修改以适应你的项目需求，欢迎下载和使用。

复制代码

import xml.etree.ElementTree as ET
from datetime import datetime
import os, cv2, time, shutil


initial_dir_orin = "E:\\lindsay\\voc_dataset_kitchen"
now = datetime.now()
year = str(now.year)
month = str(now.month).zfill(2)
day = str(now.day).zfill(2)
time_path = year+month+day
initial_dir_orin_flag = time_path+"_"+str(time.time())
annotations_path = initial_dir_orin+"/annotations"
images_path = initial_dir_orin+"/images"

new_images_path = time_path+"_repaire_datasets"+"/images"
new_annotations_path = time_path+"_repaire_datasets"+"/annotations"

if not os.path.exists(new_images_path):
    os.makedirs(new_images_path)


if not os.path.exists(new_annotations_path):
    os.makedirs(new_annotations_path)
all_labels = []
annotation_list = os.listdir(annotations_path)
print(len(annotation_list))
count = 0
# 定义一个创建一级分支object的函数
def create_object(root, xi, yi, xa, ya, obj_name):  # 参数依次，树根，xmin，ymin，xmax，ymax
    # 创建一级分支object
    _object = ET.SubElement(root, 'object')
    # 创建二级分支
    name = ET.SubElement(_object, 'name')
    # print(obj_name)
    name.text = str(obj_name)
    pose = ET.SubElement(_object, 'pose')
    pose.text = 'Unspecified'
    truncated = ET.SubElement(_object, 'truncated')
    truncated.text = '0'
    difficult = ET.SubElement(_object, 'difficult')
    difficult.text = '0'
    # 创建bndbox
    bndbox = ET.SubElement(_object, 'bndbox')
    xmin = ET.SubElement(bndbox, 'xmin')
    xmin.text = '%s' % xi
    ymin = ET.SubElement(bndbox, 'ymin')
    ymin.text = '%s' % yi
    xmax = ET.SubElement(bndbox, 'xmax')
    xmax.text = '%s' % xa
    ymax = ET.SubElement(bndbox, 'ymax')
    ymax.text = '%s' % ya


# 创建xml文件的函数
def create_tree(sources, image_name, h, w):
    imgdir = sources.split('/')[-1]
    # 创建树根annotation
    annotation = ET.Element('annotation')
    # 创建一级分支folder
    folder = ET.SubElement(annotation, 'folder')
    # 添加folder标签内容
    folder.text = (imgdir)

    # 创建一级分支filename
    filename = ET.SubElement(annotation, 'filename')
    filename.text = image_name

    # 创建一级分支path
    path = ET.SubElement(annotation, 'path')

    path.text = '{}/{}'.format(sources, image_name)  # 用于返回当前工作目录

    # 创建一级分支source
    source = ET.SubElement(annotation, 'source')
    # 创建source下的二级分支database
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'

    # 创建一级分支size
    size = ET.SubElement(annotation, 'size')
    # 创建size下的二级分支图像的宽、高及depth
    width = ET.SubElement(size, 'width')
    width.text = str(w)
    height = ET.SubElement(size, 'height')
    height.text = str(h)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'

    # 创建一级分支segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'
    return annotation

for xml in annotation_list:
    count = count+1
    img = xml.split(".xml")[0]
    image = os.path.join(images_path,img+".jpg")
    annotation = os.path.join(annotations_path,xml)
    img1 = cv2.imread(image)
    height, width, _ = img1.shape
    new_image = os.path.join(new_images_path,initial_dir_orin_flag+"_"+str(count)+".jpg")
    new_annotation = os.path.join(new_annotations_path, initial_dir_orin_flag+"_"+str(count)+".xml")
    shutil.copyfile(image, new_image)
    annotation_create = create_tree(new_images_path,initial_dir_orin_flag+"_"+str(count)+".jpg", height, width)
    print(annotation)
    tree = ET.parse(annotation)
    root = tree.getroot()
    objects = root.findall('object')
    for i, obj in enumerate(objects):
        label = obj.find('name').text
        bb = obj.find('bndbox')
        x1 = int(bb.find('xmin').text)
        y1 = int(bb.find('ymin').text)
        x2 = int(bb.find('xmax').text)
        y2 = int(bb.find('ymax').text)
        print(x1, y1, x2, y2, label, end = '\t|\t')
        all_labels.append(label)
        create_object(annotation_create, x1, y1, x2, y2, label)
    tree = ET.ElementTree(annotation_create)
    tree.write(new_annotation)
    print("\n")
print(count)

labels_list = list(set(all_labels))
labels_len = len(labels_list)
labels_str = f"this datasets have {labels_len} classes, labels list is {labels_list}"
print(labels_str)
with open("record.txt", "w") as f:
    f.write(str(count))
    f.write("\n")
    f.write(labels_str)

和任何自动化处理一样，使用这个脚本时，你应该先在少量文件上测试来确保它按预期执行，以避免意外的数据丢失或损坏。希望这个脚本和博客对你的项目有所帮助！