python：将 NWPU_VHR-10 遥感目标检测数据集转换成 YOLO 格式

在进行遥感目标检测时，NWPU_VHR-10 数据集是一个常用的高分辨率遥感图像数据集。它包含多种目标，如建筑物、飞机、车辆等。为了应用现代的目标检测算法（如 YOLO），我们需要将该数据集转换成 YOLO 所需的格式。YOLO 格式的标注通常包括图像文件和对应的标注文件，标注文件中包含目标类别及其在图像中的位置信息。

1. 数据集概述

NWPU_VHR-10 数据集由 10 种类别的遥感图像构成，每个类别的图像都有相应的标注信息，包括目标的位置和类别。在开始转换之前，首先需要理解 YOLO 格式的要求：

• 每个图像对应一个同名的 .txt 文本文件。
• 文本文件格式为：<object-class> <x_center> <y_center> <width> <height>，其中所有数值均为相对值，且范围在 [0, 1] 之间。
• x_center, y_center, width, height 的计算需基于图像的实际尺寸，以便得到相对于图像的比例。

2. 数据转换步骤

以下是将 NWPU_VHR-10 数据集转换为 YOLO 格式的基本步骤：

1. 读取原始标注文件
2. 计算目标的中心点、宽度和高度
3. 保存为 YOLO 格式

下面将提供一个示例代码，用于完成上述转换。

3. 示例代码

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

# 定义类别
classes = ['airplane', 'baseball_diamond', 'basketball_court', 'ground_track_field',
           'harbor', 'vehicle', 'swimming_pool', 'soccer_ball_field', 'tennis_court', 'building']

def convert_annotation(xml_file, output_dir, img_width, img_height):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 打开对应的文本文件
    base_name = os.path.splitext(os.path.basename(xml_file))[0]
    txt_file = os.path.join(output_dir, base_name + '.txt')

    with open(txt_file, 'w') as out_file:
        for obj in root.iter('object'):
            class_name = obj.find('name').text
            if class_name not in classes:
                continue

            class_id = classes.index(class_name)
            xmlbox = obj.find('bndbox')
            xmin = int(xmlbox.find('xmin').text)
            xmax = int(xmlbox.find('xmax').text)
            ymin = int(xmlbox.find('ymin').text)
            ymax = int(xmlbox.find('ymax').text)

            # 计算中心点和宽高
            x_center = (xmin + xmax) / 2.0 / img_width
            y_center = (ymin + ymax) / 2.0 / img_height
            width = (xmax - xmin) / img_width
            height = (ymax - ymin) / img_height

            # 写入 YOLO 格式
            out_file.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

def convert_dataset(input_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for xml_file in os.listdir(input_dir):
        if not xml_file.endswith('.xml'):
            continue

        xml_path = os.path.join(input_dir, xml_file)
        img_file = xml_file.replace('.xml', '.jpg')  # 假设图像为 JPEG 格式
        img_path = os.path.join(input_dir, img_file)

        # 读取图像以获得尺寸
        img = cv2.imread(img_path)
        img_height, img_width, _ = img.shape

        convert_annotation(xml_path, output_dir, img_width, img_height)

# 使用示例
input_directory = '/path/to/nwpu_vhr10/annotations/'
output_directory = '/path/to/yolo_format/'
convert_dataset(input_directory, output_directory)

4. 代码解释

1. 定义类别：使用了一个列表来存储可能的目标类别。
2. 读取 XML 文件：利用 xml.etree.ElementTree 从 XML 文件中读取目标信息。
3. 计算位置数据：计算出每个目标框的中心点和宽高，并将它们转换为相对值。
4. 输出结果：将结果写入到指定的输出目录。

5. 小结

通过上述步骤，我们可以将 NWPU_VHR-10 数据集转换为 YOLO 格式，以便于后续的目标检测任务。这种格式的转换使得使用 YOLO 模型进行训练和评估更加方便。希望这个示例能帮助您快速上手数据集转换的工作。