在RarePlanes数据集上怎么使用YOLOv5

发布时间：2022-03-22 10:23:06 作者：iii
来源：亿速云阅读：502

# 在RarePlanes数据集上怎么使用YOLOv5

## 引言

目标检测是计算机视觉领域的重要任务之一，而YOLOv5作为YOLO系列的最新版本之一，因其高效性和易用性受到广泛关注。RarePlanes数据集是一个专注于航空影像中飞机检测的公开数据集，包含合成和真实图像，适用于军事和民用领域的应用。本文将详细介绍如何在RarePlanes数据集上使用YOLOv5进行目标检测，包括数据准备、模型训练、评估和推理等步骤。

## 1. 环境准备

在开始之前，需要确保你的系统满足以下要求：

- Python 3.8或更高版本
- PyTorch 1.7或更高版本
- CUDA（如果使用GPU加速）
- 其他依赖库（如OpenCV、matplotlib等）

### 1.1 安装YOLOv5

首先，克隆YOLOv5的官方仓库并安装依赖：

```bash
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

2. RarePlanes数据集介绍

RarePlanes数据集由CosmiQ Works和.Reverie联合发布，包含以下特点： - 合成数据：通过模拟生成的航空影像，标注精确。 - 真实数据：来自卫星或航空摄影的真实图像。 - 标注信息：包括飞机的边界框、类别（如战斗机、客机等）和属性（如机翼形状、发动机数量等）。

数据集通常以COCO或Pascal VOC格式提供，需要转换为YOLOv5支持的格式。

2.1 数据集下载

从官方渠道下载RarePlanes数据集： - 访问RarePlanes官网或.Reverie平台。 - 下载数据集并解压到本地目录，例如data/rareplanes。

2.2 数据集结构

解压后的数据集结构如下：

rareplanes/
├── train/
│   ├── images/
│   └── labels/
├── val/
│   ├── images/
│   └── labels/
└── test/
    ├── images/
    └── labels/

2.3 数据格式转换

如果数据集是COCO格式，需要转换为YOLOv5的TXT格式（每行一个标注，格式为class_id x_center y_center width height）。可以使用以下脚本：

import json
import os

def coco_to_yolo(coco_json, output_dir):
    with open(coco_json) as f:
        data = json.load(f)
    
    os.makedirs(output_dir, exist_ok=True)
    
    for image in data['images']:
        image_id = image['id']
        file_name = image['file_name']
        width = image['width']
        height = image['height']
        
        annotations = [ann for ann in data['annotations'] if ann['image_id'] == image_id]
        
        with open(os.path.join(output_dir, file_name.replace('.png', '.txt')), 'w') as f:
            for ann in annotations:
                class_id = ann['category_id']
                x, y, w, h = ann['bbox']
                x_center = x + w / 2
                y_center = y + h / 2
                f.write(f"{class_id} {x_center/width} {y_center/height} {w/width} {h/height}\n")

# 示例调用
coco_to_yolo('rareplanes/annotations/train.json', 'rareplanes/train/labels')

3. 配置YOLOv5

3.1 创建数据集配置文件

在yolov5/data/目录下创建rareplanes.yaml文件，内容如下：

# rareplanes.yaml
train: ../data/rareplanes/train/images
val: ../data/rareplanes/val/images
test: ../data/rareplanes/test/images

# 类别数量和名称
nc: 5  # 根据实际类别数调整
names: ['fighter', 'passenger', 'cargo', 'helicopter', 'other']  # 示例类别

3.2 选择模型

YOLOv5提供多种预训练模型（如YOLOv5s、YOLOv5m、YOLOv5l等）。根据需求选择： - YOLOv5s：轻量级，适合快速实验。 - YOLOv5l：更大模型，精度更高。

4. 训练模型

使用以下命令开始训练：

python train.py --img 640 --batch 16 --epochs 100 --data data/rareplanes.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt --name rareplanes_exp

参数说明： - --img 640：输入图像尺寸。 - --batch 16：批次大小（根据GPU显存调整）。 - --epochs 100：训练轮数。 - --data：数据集配置文件路径。 - --cfg：模型配置文件路径。 - --weights：预训练权重路径。 - --name：实验名称。

4.1 训练监控

YOLOv5会自动记录训练日志，并通过TensorBoard可视化：

tensorboard --logdir runs/train

5. 模型评估

训练完成后，使用以下命令评估模型性能：

python val.py --data data/rareplanes.yaml --weights runs/train/rareplanes_exp/weights/best.pt --img 640

输出包括mAP、精确率、召回率等指标。

6. 模型推理

使用训练好的模型进行预测：

python detect.py --source ../data/rareplanes/test/images --weights runs/train/rareplanes_exp/weights/best.pt --conf 0.25 --img 640

结果会保存在runs/detect/目录下。

7. 模型优化

7.1 数据增强

在rareplanes.yaml中配置数据增强参数：

# 数据增强
hsv_h: 0.015  # 色调增强
hsv_s: 0.7    # 饱和度增强
hsv_v: 0.4    # 明度增强
flipud: 0.5   # 上下翻转概率

7.2 超参数调优

使用hyp.scratch.yaml调整学习率、权重衰减等超参数。

8. 部署模型

将模型导出为ONNX或TorchScript格式以便部署：

python export.py --weights runs/train/rareplanes_exp/weights/best.pt --include onnx

9. 常见问题

9.1 显存不足

减小--batch大小。
使用--img 320降低输入尺寸。

9.2 标注错误

检查数据格式是否正确。
使用LabelImg等工具可视化标注。

10. 结论

本文详细介绍了在RarePlanes数据集上使用YOLOv5进行目标检测的全流程。通过合理配置和调优，YOLOv5可以在航空影像中实现高效的飞机检测。未来可以尝试： - 结合更多数据增强提升泛化性。 - 使用更大的模型（如YOLOv5x）提高精度。 - 部署到边缘设备（如无人机）实现实时检测。

参考资料

YOLOv5官方仓库：https://github.com/ultralytics/yolov5
RarePlanes数据集官网：https://www.cosmiqworks.org/RarePlanes/
COCO数据格式说明：https://cocodataset.org/

”`