代码收藏家 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幅经锐化的彩色红外图像。
该数据集被分成两组:
由此，正面图像集、757飞机、302船只、655储罐、390棒球钻石、524网球场、159篮球场、163地面田径场、224港口、124桥梁和477车辆被手动注释为边界框和用于地面真相的实例面具。
使用全部数据集或者部分数据的时候，需要引用以下论文：
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的标注格式非常简洁，对于每张图片中的每一个对象，都会有一行描述该对象的信息。每一行包含五个值：
举个例子，如果有一个图像，其中包含一个类别为“人”的对象，该对象的边界框中心点位于图像的 (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节。
作者：_养乐多_
