NWPU VHR-10 数据集是地理空间目标检测领域的经典挑战集，包含 650 张高分辨率遥感图像。数据源自 Google Earth 等平台，覆盖飞机、船舶、储罐等 10 类常见目标，非常适合用于验证小目标检测算法的鲁棒性。
数据集概览
该数据集共标注了 3,896 个目标实例，按默认比例划分为训练集（390 张）、验证集（130 张）和测试集（130 张）。所有目标均采用水平边界框（Axis-Aligned Bounding Box）标注，支持 YOLO TXT 和 COCO JSON 两种主流格式。
| 类别 ID | 英文名称 | 中文名称 | 数量 |
|---|---|---|---|
| 0 | Airplane | 飞机 | 757 |
| 1 | Ship | 船只 | 302 |
| 2 | Storage Tank | 储罐 | 655 |
| 3 | Baseball Diamond | 棒球场 | 390 |
| 4 | Tennis Court | 网球场 | 524 |
| 5 | Basketball Court | 篮球场 | 159 |
| 6 | Ground Track Field | 田径场 | 163 |
| 7 | Harbor | 港口 | 224 |
| 8 | Bridge | 桥梁 | 124 |
| 9 | Vehicle | 车辆 | 598 |
注意：遥感场景下背景复杂，且存在明显的尺度差异（如港口与车辆），训练时需重点关注数据增强策略。
项目结构搭建
建议将数据整理为如下目录结构，便于后续脚本调用：
nwpu_vhr10_yolo/
├── dataset.yaml
├── train.py
├── predict_demo.py
├── export_model.py
└── data/
├── images/
│ ├── train/
│ ├── val/
│ └── test/
└── labels/
├── train/
├── val/
└── test/
假设你已将官方提供的 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 加速。
