【Python快速入门和实践011】Python常用脚本-目标检测之VOC格式转YOLO格式脚本

一、数据集介绍

        NEU-DET 数据集是由东北大学（Northeastern University, 简称 NEU）发布的一个用于钢材表面缺陷检测的数据集。这个数据集特别设计用于支持和促进工业领域中的缺陷检测研究。

NEU-DET 数据集的一些主要特点包括：

• 多样性和复杂性：数据集包含了多种类型的钢材表面缺陷，如裂纹、划痕、凹坑、斑点等。
• 高质量标注：每张图像都经过仔细标注，包括缺陷的位置、大小以及类别等详细信息。
• 应用场景：主要用于训练和评估计算机视觉模型，特别是用于缺陷检测的任务，可以用于训练目标检测模型。
• 规模：虽然具体的数据集大小没有详细说明，但作为一个专业研究数据集，它应该包含足够的图像样本以支持深度学习模型的训练。

二、VOC格式介绍

        VOC（Visual Object Classes）格式是一种广泛使用的图像注释标准，常用于计算机视觉任务，如物体检测、分类和分割。VOC 数据集最初是由 Pascal Visual Object Classes Challenge 发起的，该挑战赛自 2005 年开始举办，旨在推动计算机视觉技术的发展。

VOC 数据集的特点

• 标准化：VOC 格式定义了一套标准化的数据结构和文件组织方式，使得不同的研究人员能够方便地共享和使用相同的数据集。
• 多样化：VOC 数据集通常包含多个类别的物体，如人、车、动物等。
• 详细标注：每个图像都附有详细的标注信息，包括物体的位置（用边界框表示）、类别标签、分割掩码等。
• 评估工具：VOC 数据集提供了用于评估检测算法性能的工具，包括计算平均精度（Average Precision, AP）等指标的方法。

数据集组成

• 图像：每个数据集中包含一系列 JPEG 图像。
• 标注文件：对于每个图像，有一个 XML 文件存储了图像中所有物体的标注信息。
  • XML 文件结构：
    • <annotation>：根节点。
    • <filename>：图像文件名。
    • <folder>：图像所在的文件夹。
    • <source>：数据源信息。
    • <size>：图像尺寸信息，包括宽度、高度和通道数。
    • <segmented>：是否进行了分割标注。
    • <object>：每个物体的标注信息。
      • <name>：物体类别名称。
      • <pose>：物体的姿态（可选）。
      • <truncated>：物体是否被截断。
      • <difficult>：物体是否难以检测。
      • <bndbox>：物体的边界框坐标。
        • <xmin>、<ymin>、<xmax>、<ymax>：边界框的四个角坐标。

VOC格式如下：

<?xml version="1.0" encoding="utf-8"?>
	<annotation>
		<folder>driving_annotation_dataset</folder>
		<filename>crazing_1.jpg</filename>
		<size>
			<width>200</width>
			<height>200</height>
			<depth>3</depth>
		</size>
		<object>
			<name>crazing</name>
			<pose>Unspecified</pose>
			<truncated>0</truncated>
			<difficult>0</difficult>
			<bndbox>
				<xmin>3</xmin>
				<ymin>3</ymin>
				<xmax>194</xmax>
				<ymax>195</ymax>
			</bndbox>
		</object>
	</annotation>

三、YOLO格式介绍

        YOLO (You Only Look Once) 格式是一种广泛应用于目标检测任务的数据标注格式。这种格式非常简洁且易于处理，非常适合用于训练和评估基于 YOLO 的目标检测模型。

YOLO 格式的特点

• 简洁性：YOLO 格式非常简单，每个标注文件只包含一个文本文件，其中包含图像中每个目标对象的标注信息。
• 标准化：YOLO 格式定义了一个统一的数据结构，便于数据的处理和模型的训练。
• 高效性：由于其简洁性，YOLO 格式非常适合大规模数据集的处理和模型训练。

数据集组成

• 图像：每个数据集包含一系列 JPEG 或 PNG 格式的图像文件。
• 标注文件：对于每个图像，都有一个对应的文本文件（通常是 .txt 扩展名）来存储标注信息。

标注文件结构

每个标注文件（.txt 文件）包含一行或多行文本，每一行对应图像中的一个目标对象。每行文本包含以下信息：

• 类别编号：目标对象所属类别的编号（从 0 开始）。
• 中心点 x 坐标：目标对象中心点的 x 坐标，相对于图像宽度归一化。
• 中心点 y 坐标：目标对象中心点的 y 坐标，相对于图像高度归一化。
• 宽度：目标对象的宽度，相对于图像宽度归一化。
• 高度：目标对象的高度，相对于图像高度归一化。

YOLO格式如下：

0 0.45 0.50 0.35 0.40
1 0.80 0.75 0.10 0.15

这表示：

• 第一个对象属于类别 0，在图像中的位置为 (0.45, 0.50)，宽度为 0.35，高度为 0.40。
• 第二个对象属于类别 1，在图像中的位置为 (0.80, 0.75)，宽度为 0.10，高度为 0.15。

使用场景

• 训练模型：用于训练 YOLO 及其变体的目标检测模型。
• 评估性能：用于评估模型的性能，通过计算平均精度（mAP）等指标来衡量模型的准确性。

获取方式

• 创建自己的数据集：可以自己收集图像并使用相应的工具对其进行标注，生成 YOLO 格式的标注文件。
• 现成的数据集：许多开源数据集已经提供了 YOLO 格式的标注文件，可以在 GitHub 或其他数据集共享平台上找到它们。

四、VOC转YOLO格式脚本

import xml.etree.ElementTree as ET
import os


# 定义函数将坐标转换为YOLO格式
def convert(size, box):
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]
    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]
    return (x, y, w, h)


# 定义函数将VOC格式的XML标签文件转换为YOLO格式的TXT标签文件
def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    for xml_name in xml_files:
        xml_file = os.path.join(xml_files_path, xml_name)
        base_name = os.path.splitext(xml_name)[0]  # 获取文件名的基本名称
        out_txt_path = os.path.join(save_txt_files_path, base_name + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                continue
            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)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


# 主程序入口
if __name__ == "__main__":
    # 1、需要转化的类别
    classes = [ 'rolled-in_scale', 'patches','crazing','pitted_surface','inclusion','scratches' ]  # 根据自己的类别名称及种类自行更改
    # 2、VOC格式的XML标签文件路径
    xml_files1 = r'E:\pythonProject\pythonProject\NEU-DET\xmls'
    # 3、转化为YOLO格式的TXT标签文件存储路径
    save_txt_files1 = r'E:\pythonProject\pythonProject\NEU-DET\txt'

    convert_annotation(xml_files1, save_txt_files1, classes)

转换代码只需要改变你数据集的类别，然后指定xml和保存yolo格式的txt文件夹路径即可。