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

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

配置文件编写

在根目录下创建 dataset.yaml，定义路径及类别名称。这里以 YOLOv8 为例：

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

模型训练实战

接下来我们处理训练环节。由于遥感图像中小目标较多，推荐使用中等规模模型（如 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,
        batch=16,
        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,
        translate=0.1,
        scale=0.5,
        flipud=0.0,
        fliplr=0.5,
        mosaic=1.0,
        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)

优化建议

在实际应用中，可能会遇到以下挑战，建议针对性调整：

• 小目标漏检：启用 mosaic=1.0，若显存允许可将输入尺寸提升至 1024。
• 类别不平衡：桥梁等样本较少，YOLOv8 内置 Focal Loss 通常能缓解，也可尝试对稀有类进行过采样。
• 背景干扰：增加 HSV 变换范围和旋转角度，提升光照与视角的鲁棒性。
• 实时性要求：边缘端部署建议改用 yolov8s 并结合 TensorRT 加速。