NWPU VHR-10 遥感目标检测实战指南

所有目标均以水平边界框（Axis-Aligned Bounding Box）标注。

环境准备与配置

建议按以下结构组织项目文件，确保路径引用准确：

nwpu_vhr10_yolo/
├── dataset.yaml
├── train.py
├── predict_demo.py
├── export_model.py
└── data/
    ├── images/
    │   ├── train/
    │   ├── val/
    │   └── test/
    └── labels/
        ├── train/
        ├── val/
        └── test/

创建 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']

模型训练实战

使用 Ultralytics YOLOv8 进行训练。由于遥感图像中小目标较多，建议选择中等规模模型如 yolov8m，并开启 Mosaic 增强以提升鲁棒性。

# 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,
        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

推理演示

训练完成后，加载最佳权重进行单图预测。注意遥感小目标建议适当降低置信度阈值。

# 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})")

模型导出与部署

为了在实际场景中应用，可以将模型导出为 ONNX、TensorRT 或 OpenVINO 格式。

# export_model.py
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 的特性，以下是几个关键优化点：

• 小目标漏检：启用 mosaic=1.0，若显存允许将 imgsz 提升至 1024。
• 类别不平衡：桥梁类样本较少（124 个），YOLOv8 内置 Focal Loss 变体通常能处理，也可尝试对稀有类过采样。
• 背景复杂：增加 hsv_v 和 degrees 增强参数，提升光照与角度变化的鲁棒性。
• 边缘部署：若运行于无人机等设备，推荐使用 yolov8s + TensorRT 加速。