如何使用NWPU VHR-10数据集训练YOLO模型：遥感目标检测实战教程

如何使用NWPU VHR-10数据集训练YOLO模型：遥感目标检测实战教程

【免费下载链接】NWPU-VHR-10dataset 项目地址: https://ai.gitcode.com/GewisLab/NWPU-VHR-10dataset

NWPU VHR-10数据集是西北工业大学发布的10类遥感图像目标检测数据集，包含800张高分辨率遥感图像，适用于训练YOLO等目标检测模型。本教程将详细介绍如何利用该数据集构建遥感目标检测系统，从数据准备到模型训练的完整流程。

1. 数据集核心特性解析

NWPU VHR-10数据集专为地理空间目标检测设计，包含10个常见遥感目标类别：

编号	类别 (英文)	类别 (中文)
1	Airplane	飞机
2	Ship	船舶
3	Storage tank	储油罐
4	Baseball diamond	棒球场
5	Tennis court	网球场
6	Basketball court	篮球场
7	Ground track field	田径场
8	Harbor	港口
9	Bridge	桥梁
10	Vehicle	车辆

数据集文件结构清晰，分为三个主要部分：

• negative image set/：150张不含任何目标的背景图像
• positive image set/：650张包含至少一个目标的图像
• ground truth/：650个文本文件，对应positive image set中的图像标注信息

2. 数据准备：从下载到格式转换

2.1 数据集获取

首先克隆数据集仓库：

git clone https://gitcode.com/GewisLab/NWPU-VHR-10dataset
cd NWPU-VHR-10dataset

2.2 标注格式解析

ground truth目录下的文本文件采用以下格式标注目标：

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

• (x1,y1)：边界框左上角坐标
• (x2,y2)：边界框右下角坐标
• a：目标类别编号（1-10，对应上述类别表）

例如，标注文件ground truth/001.txt可能包含：

(100,200),(300,400),1
(500,600),(700,800),2

表示图像中包含一个飞机（类别1）和一个船舶（类别2）。

2.3 转换为YOLO格式

YOLO模型需要特定格式的标注文件（每行一个目标）：

<class_id> <x_center> <y_center> <width> <height>

其中坐标需归一化到0-1范围。以下Python代码可实现格式转换：

import os
import cv2

def convert_nwpu_to_yolo(nwpu_root, output_dir):
    # 创建输出目录
    os.makedirs(os.path.join(output_dir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True)
    
    # 处理所有正样本图像
    for img_name in os.listdir(os.path.join(nwpu_root, 'positive image set')):
        if not img_name.endswith('.jpg'):
            continue
            
        # 读取图像获取尺寸
        img_path = os.path.join(nwpu_root, 'positive image set', img_name)
        img = cv2.imread(img_path)
        if img is None:
            continue
        h, w = img.shape[:2]
        
        # 读取对应标注文件
        label_name = os.path.splitext(img_name)[0] + '.txt'
        label_path = os.path.join(nwpu_root, 'ground truth', label_name)
        
        if not os.path.exists(label_path):
            continue
            
        # 转换标注格式
        yolo_labels = []
        with open(label_path, 'r') as f:
            for line in f.readlines():
                line = line.strip()
                if not line:
                    continue
                    
                # 解析NWPU格式
                coords, cls = line.rsplit(',', 1)
                (x1, y1), (x2, y2) = eval(coords)
                
                # 转换为YOLO格式
                x_center = (x1 + x2) / 2 / w
                y_center = (y1 + y2) / 2 / h
                width = (x2 - x1) / w
                height = (y2 - y1) / h
                
                # 类别ID减1（YOLO从0开始）
                yolo_labels.append(f"{int(cls)-1} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
        
        # 保存转换结果
        if yolo_labels:
            # 复制图像
            dst_img_path = os.path.join(output_dir, 'images', img_name)
            cv2.imwrite(dst_img_path, img)
            
            # 保存YOLO标注
            dst_label_path = os.path.join(output_dir, 'labels', label_name)
            with open(dst_label_path, 'w') as f:
                f.write('\n'.join(yolo_labels))

# 使用示例
convert_nwpu_to_yolo('./', './yolo_dataset')

3. 配置YOLO模型训练环境

3.1 安装必要依赖

# 创建虚拟环境
conda create -n yolo-env python=3.8 -y
conda activate yolo-env

# 安装PyTorch (根据CUDA版本调整)
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

# 安装YOLOv5
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

3.2 创建数据集配置文件

在yolov5/data目录下创建nwpu_vhr10.yaml：

# NWPU VHR-10数据集配置
path: ../yolo_dataset  # 数据集根目录
train: images  # 训练集图像目录
val: images  # 验证集图像目录（实际应用中应划分单独验证集）

# 类别信息
nc: 10  # 类别数量
names: ['Airplane', 'Ship', 'Storage tank', 'Baseball diamond', 'Tennis court', 
        'Basketball court', 'Ground track field', 'Harbor', 'Bridge', 'Vehicle']

4. 模型训练与评估

4.1 启动训练

# 使用预训练模型开始训练
python train.py --img 640 --batch 16 --epochs 100 --data nwpu_vhr10.yaml --weights yolov5s.pt --cache

关键参数说明：

• --img 640：输入图像尺寸
• --batch 16：批次大小（根据GPU内存调整）
• --epochs 100：训练轮数
• --data：数据集配置文件
• --weights：预训练权重

4.2 训练过程监控

训练过程中可通过TensorBoard监控指标：

tensorboard --logdir runs/train

主要关注指标：

• mAP@0.5：IoU=0.5时的平均精度
• loss：训练损失
• precision/recall：精确率和召回率

4.3 模型评估

训练完成后，使用验证集评估模型性能：

python val.py --weights runs/train/exp/weights/best.pt --data nwpu_vhr10.yaml --img 640

5. 模型推理与应用

5.1 单张图像推理

python detect.py --weights runs/train/exp/weights/best.pt --source ../NWPU-VHR-10dataset/positive\ image\ set/001.jpg --conf 0.5

5.2 批量处理图像

python detect.py --weights runs/train/exp/weights/best.pt --source ../NWPU-VHR-10dataset/positive\ image\ set/ --conf 0.5 --save-txt

6. 优化建议与常见问题

6.1 提升模型性能的技巧

1. 数据增强：在yolov5/data/hyps/hyp.scratch-low.yaml中调整增强参数
2. 学习率调度：尝试不同的学习率策略
3. 模型选择：对于小目标检测，可尝试YOLOv5s或YOLOv5m模型
4. 多尺度训练：添加--multi-scale参数进行多尺度训练

6.2 常见问题解决

• 标注错误：检查ground truth文件格式，确保坐标正确
• 类别不平衡：可使用--weighted-loss参数处理类别不平衡问题
• 内存不足：减小批次大小或图像尺寸
• 过拟合：增加数据增强，使用早停策略

7. 数据集引用与学术使用

使用该数据集发表成果时，请引用以下论文：

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.

⚠️ 注意：本数据集仅用于研究目的，原始数据集版权归西北工业大学及其作者所有。

【免费下载链接】NWPU-VHR-10dataset 项目地址: https://ai.gitcode.com/GewisLab/NWPU-VHR-10dataset