NWPU VHR-10 遥感目标检测数据集：从数据解析到 YOLO 实战

注意：类间尺度差异较大（如港口与车辆），且存在较多小目标，这对模型泛化能力提出了较高要求。

二、项目结构搭建

为了便于管理，建议采用如下目录结构组织数据与代码：

nwpu_vhr10_yolo/
├── dataset.yaml      # 数据集配置文件
├── train.py          # 主训练脚本
├── predict_demo.py   # 推理演示
├── export_model.py   # 模型导出脚本
└── data/
    ├── images/
    │   ├── train/    # 390 张训练图像
    │   ├── val/      # 130 张验证图像
    │   └── test/     # 130 张测试图像
    └── labels/
        ├── train/    # YOLO 格式标注
        ├── val/
        └── test/

三、配置与训练

1. 数据集配置 (dataset.yaml)

在 dataset.yaml 中定义路径和类别名称，确保与目录结构一致。

path: ./data
train: images/train
val: images/val
test: images/test
nc: 10
names:
  - airplane
  - ship
  - storage_tank
  - baseball_diamond
  - tennis_court
  - basketball_court
  - ground_track_field
  - harbor
  - bridge
  - vehicle

2. 训练脚本 (train.py)

针对遥感图像特点，我们在训练参数上做了一些针对性调整，特别是针对小目标和多尺度的增强策略。

from ultralytics import YOLO
import torch

def main():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"🚀 使用设备：{device}")

    # 选择模型，遥感小目标多，推荐 yolov8m 或 yolov8l
    model_name = "yolov8m.pt"
    project_name = "nwpu_vhr10_detection"

    # 加载预训练模型
    model = YOLO(model_name)

    # 开始训练
    results = model.train(
        data='dataset.yaml',
        epochs=150,           # 遥感数据复杂，建议充分训练
        imgsz=640,            # 640 足够；若需更高精度可用 1024
        batch=16,             # 显存允许可设为 32
        name=project_name,
        device=device,
        patience=30,          # 早停机制
        save=True,
        plots=True,           # 自动生成 PR 曲线、混淆矩阵
        exist_ok=False,
        workers=8,
        optimizer='AdamW',
        lr0=0.01,
        lrf=0.01,
        augment=True,         # 关键！提升鲁棒性
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        degrees=15.0,         # ±15° 旋转（适应不同视角）
        translate=0.1,
        scale=0.5,
        flipud=0.0,           # 不上下翻转（地理目标方向敏感）
        fliplr=0.5,           # 左右翻转安全
        mosaic=1.0,           # Mosaic 增强对小目标极有效
        mixup=0.1
    )

    # 在测试集上评估最终性能
    metrics = model.val(data='dataset.yaml', split='test')
    print("\n✅ NWPU VHR-10 测试集结果:")
    print(f" [email protected] (all): {metrics.box.map50:.5f}")
    print(f" [email protected]:0.95: {metrics.box.map:.5f}")
    print(f" Precision: {metrics.box.mp:.5f}")
    print(f" Recall: {metrics.box.mr:.5f}")

if __name__ == '__main__':
    main()

运行命令：

python train.py

四、推理与部署

1. 推理演示 (predict_demo.py)

训练完成后，我们可以加载权重进行单图或视频推理。对于遥感小目标，适当降低置信度阈值有助于提高召回率。

from ultralytics import YOLO

model = YOLO('runs/detect/nwpu_vhr10_detection/weights/best.pt')
results = model.predict(
    source='test_image.jpg',
    conf=0.25,              # 遥感小目标建议降低阈值
    iou=0.45,
    save=True,
    show=True,
    line_width=2,
    font_size=12
)

# 打印检测结果
class_names = ['Airplane','Ship','Storage Tank','Baseball Diamond','Tennis Court','Basketball Court','Ground Track Field','Harbor','Bridge','Vehicle']
for result in results:
    boxes = result.boxes
    print(f"检测到 {len(boxes)} 个目标")
    for box in boxes:
        cls_id = int(box.cls.item())
        conf = float(box.conf.item())
        print(f" → {class_names[cls_id]} (置信度：{conf:.3f})")

2. 模型导出

为了将模型部署到边缘设备或生产环境，可以导出为 ONNX、TensorRT 或 OpenVINO 格式。

from ultralytics import YOLO

model = YOLO('runs/detect/nwpu_vhr10_detection/weights/best.pt')

# 导出为 ONNX（通用）
model.export(format='onnx', imgsz=640, simplify=True)

# 导出为 TensorRT（NVIDIA GPU / Jetson）
model.export(format='engine', imgsz=640, half=True, device=0)

# 导出为 OpenVINO（Intel CPU）
model.export(format='openvino', imgsz=640)

五、优化建议

在实际项目中，可能会遇到以下问题，可以参考对应策略：

通过上述流程，你可以快速完成 NWPU VHR-10 数据集的接入与模型训练，并根据实际需求进行优化部署。