NWPU VHR-10 遥感目标检测数据集详解与 YOLO 训练实战

🔢 总标注目标数：3,896 个
📌 所有目标均以 水平边界框（Axis-Aligned Bounding Box） 标注。

三、推荐使用方式

四、项目结构建议

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

✅ 假设你已将官方提供的 YOLO 格式数据按上述结构组织。

五、数据集配置文件 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']

六、详细训练代码 train.py

# 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"
    # 可选：yolov8s / yolov8l / yolov8x
    project_name = "nwpu_vhr10_detection"
    
    # 加载预训练模型
    model = YOLO(model_name)
    
    # 开始训练
    results = model.train(
        data='dataset.yaml',
        epochs=150,  # 遥感数据复杂，建议充分训练
        imgsz=640,  # NWPU 图像较小，640 足够；若需更高精度可用 1024
        batch=16,  # 640 分辨率下 batch=16（32G 显存可设 32）
        name=project_name,
        device=device,
        patience=30,  # 早停机制（验证集 mAP 不升则停）
        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 增强（对小目标如 vehicle/airplane 极有效）
        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

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

八、模型导出（用于部署）

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

九、训练优化建议