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

作者：CSDN @ _养乐多_

本文将介绍将 NWPU_VHR-10 遥感目标检测数据集转换成 YOLO 格式的 python 脚本。

---

---

一、数据集介绍

1.1 数据集下载

https://opendatalab.com/OpenDataLab/NWPU_VHR-10

1.2 数据集介绍

NWPU VHR-10数据集是具有挑战性的十类地理空间对象检测数据集。该数据集总共包含800幅VHR光学遥感图像，其中715幅彩色图像是从谷歌地球获得的，空间分辨率从0.5到2米，从具有0.08 m的空间分辨率的Vaihingen数据获取85幅经锐化的彩色红外图像。

该数据集被分成两组:

• a) 在图像中包含至少一个目标的正图像集包含650图像。
• b) 负图像集包含150图像，并且它不包含任何目标。
  由此，正面图像集、757飞机、302船只、655储罐、390棒球钻石、524网球场、159篮球场、163地面田径场、224港口、124桥梁和477车辆被手动注释为边界框和用于地面真相的实例面具。

使用全部数据集或者部分数据的时候，需要引用以下论文：

• Gong Cheng, Junwei Han, Peicheng Zhou, Lei Guo. Multi-class geospatial object detection and geographic image classification based on collection of part detectors. ISPRS Journal of Photogrammetry and Remote Sensing, 98: 119-132, 2014.
• Gong Cheng, Junwei Han. A survey on object detection in optical remote sensing images. ISPRS Journal of Photogrammetry and Remote Sensing, 117: 11-28, 2016.
• Gong Cheng, Peicheng Zhou, Junwei Han. Learning rotation-invariant convolutional neural networks for object detection in VHR optical remote sensing images. IEEE Transactions on Geoscience and Remote Sensing, 54(12): 7405-7415, 2016.

1.3 数据格式

文本文件的每一行定义了一个对象边界框，格式如下：

(x1,y1),(x2,y2),a

其中(x1,y1)表示边界框的左上角坐标，(x2,y2)表示边界框的右下角坐标，而a是对象类别（1-飞机，2-船只，3-储罐，4-棒球场，5-网球场，6-篮球场，7-田径场地，8-港口，9-桥梁，10-车辆）。

二、格式转换

我们需要将 NWPU_VHR-10 数据集转换成 YOLO 格式，以进行目标检测。

就是将

(x1,y1),(x2,y2),a

转换为 YOLO 格式

class, x_center, y_center, width, height

YOLO的标注格式非常简洁，对于每张图片中的每一个对象，都会有一行描述该对象的信息。每一行包含五个值：

• 类别的索引（class index）：表示该对象所属的类别，在类别列表中的索引（从0开始）。
• 中心点的 x 坐标（x_center）：对象边界框中心点相对于图像宽度的归一化位置（范围 0 到 1）。
• 中心点的 y 坐标（y_center）：对象边界框中心点相对于图像高度的归一化位置（范围 0 到 1）。
• 边界框的宽度（width）：边界框的宽度相对于图像宽度的归一化值（范围 0 到 1）。
• 边界框的高度（height）：边界框的高度相对于图像高度的归一化值（范围 0 到 1）。

举个例子，如果有一个图像，其中包含一个类别为“人”的对象，该对象的边界框中心点位于图像的 (0.4, 0.5)，边界框的宽度为图像宽度的 0.1 倍，高度为图像高度的 0.2 倍，并且“人”这个类别的索引是 0，则对应的 YOLO 格式的标注文件中的一行可能如下所示：

0 0.4 0.5 0.1 0.2

三、完整代码

使用时，需要修改 class_map 和文件路径。

import os
from PIL import Image


def convert_to_yolo_format(x1, y1, x2, y2, img_width, img_height):
    dw = 1.0 / img_width
    dh = 1.0 / img_height
    center_x = (x1 + x2) / 2.0
    center_y = (y1 + y2) / 2.0
    width = x2 - x1
    height = y2 - y1
    center_x *= dw
    center_y *= dh
    width *= dw
    height *= dh
    return (center_x, center_y, width, height)


# 文件路径
annotations_dir = r'E:\DataSet\NWPUVHR10dataset\annotations'
images_dir = r'E:\DataSet\NWPUVHR10dataset\positive'
output_dir = r'E:\DataSet\NWPUVHR10dataset\yolo_labels'

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 类别映射
class_map = {
    '1': '0',  # 飞机
    '2': '1',  # 船只
    '3': '2',  # 储油罐
    '4': '3',  # 棒球场
    '5': '4',  # 网球场
    '6': '5',  # 篮球场
    '7': '6',  # 跑道场地
    '8': '7',  # 港口
    '9': '8',  # 桥梁
    '10': '9'  # 车辆
}


def get_image_size(image_path):
    # 打开图片文件
    with Image.open(image_path) as img:
        # 获取图片的宽度和高度
        width, height = img.size
        return width, height


def parse_bbox(line):
    try:
        # 去掉括号，并用逗号分割
        line = line.strip().replace('(', '').replace(')', '')
        parts = line.split(',')
        if len(parts) != 5:
            raise ValueError("Invalid line format")

        x1, y1, x2, y2 = map(int, parts[:4])
        class_id = parts[4]

        return x1, y1, x2, y2, class_id
    except ValueError as e:
        print(f"Error parsing line: {line}. Error: {e}")
        return None, None, None, None, None


# 遍历标注文件
for annotation_file in os.listdir(annotations_dir):
    if annotation_file.endswith('.txt'):
        # 读取对应的图像文件
        image_file = os.path.splitext(annotation_file)[0] + '.jpg'
        image_path = os.path.join(images_dir, image_file)

        # 打印调试信息
        print(f"Processing image: {image_path}")

        if not os.path.exists(image_path):
            print(f"Image file {image_path} does not exist, skipping.")
            continue

        # 获取图像大小
        img_width, img_height = get_image_size(image_path)

        # 读取标注文件
        annotation_path = os.path.join(annotations_dir, annotation_file)
        with open(annotation_path, 'r') as f:
            lines = f.readlines()

        # 输出文件路径
        output_path = os.path.join(output_dir, annotation_file)

        with open(output_path, 'w') as out_f:
            for line in lines:
                # 解析标注信息
                x1, y1, x2, y2, class_id = parse_bbox(line)
                if x1 is None:
                    continue

                class_id = class_map.get(class_id, None)
                if class_id is None:
                    print(f"Unknown class id: {class_id}")
                    continue

                # 转换为YOLO格式
                center_x, center_y, width, height = convert_to_yolo_format(x1, y1, x2, y2, img_width, img_height)

                # 写入输出文件
                out_f.write(f"{class_id} {center_x} {center_y} {width} {height}\n")

print("Conversion completed.")

四、划分数据集

如果需要将图片和标签数据集按比例划分为训练、验证、测试数据集，请参考以下博客。

参考博客《YOLO：VOC格式数据集转换为YOLO数据集格式》中的第2节。