【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 复制代码

<?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格式如下：

XML 复制代码

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格式脚本

python 复制代码

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文件夹路径即可。