YOLO模型

一种实时物体检测模型，速度快精度较高，适合非高精度的实时处理场景。

一.环境搭建

需要安装python3和依赖项

1.安装虚拟环境和依赖项

  • 安装python3和pip

    sudo apt update
    sudo apt install python3 python3-pip python3-venv

  • 创建虚拟环境（避免与系统安装的python库发生冲突）

    python3 -m venv yolov5_env
    source yolov5_env/bin/activate

  • 安装YOLOv5的依赖项：

    pip install -U pip # 确保 pip 是最新的
    pip install torch torchvision torchaudio # 安装 PyTorch（YOLOv5 的依赖）
    pip install matplotlib numpy opencv-python # 安装其他常用库
    pip install -U pycocotools # 安装 pycocotools（用于评估）

之后所有命令都在yolov5虚拟环境下进行！

2.克隆YOLOv5仓库并安装依赖

YOLOv5 是由 ultralytics 提供的，可以直接从 GitHub 克隆仓库并安装依赖项。

二.准备无人机数据集

需要准备包含无人机标注的图像数据集,YOLO 使用COCO格式或YOLO格式的数据集，通常包括图像文件和标签文件。

1.数据集格式

  • YOLO 格式要求每张图像对应一个 .txt 文件，每个文件包含图像中所有目标的边界框信息，格式如下：

    <class_id> <x_center> <y_center> <width> <height>

  • 所有值都需要归一化到 [0, 1] 范围内。例如：

    0 0.5 0.5 0.1 0.1

2.下载数据集（也可以直接将压缩包下载解压）

可以在Kaggle上寻找合适的数据集并下载，还有VisDrone、UAVDT、DroneNet。以Kaggle为例：

  • 获取 Kaggle API 密钥：

    登录Kaggle 后，点击右上角的头像，选择Account，在API部分，点击Create New API Token按钮，这将会下载一个名为 kaggle.json 的文件。

  • 安装Kaggle API 客户端:

    pip install kaggle

  • 配置Kaggle API 密钥,并下载的 kaggle.json 文件移动到该文件夹中

    mkdir ~/.kaggle
    mv /path/to/kaggle.json ~/.kaggle/

  • 使用 Kaggle API 下载数据集
    UR格式https://www.kaggle.com/datasets/xxxxxx/your-dataset-name中的xxxxxx是数据集的 ID

    kaggle datasets download -d

  • 解压数据集

    unzip your-dataset-name.zip -d /path/to/extract/directory

3.配置数据集路径

  • 重新组织数据集目录，结构如下：

    ~/yolov5/data/

  • 将文件移动到相应的文件夹中

  • YOLOv5 训练时需要一个配置文件，即data.yaml 文件，需要自己nano创建并编写：

    train: /path/to/your/train/images # 训练集图像路径
    val: /path/to/your/val/images # 验证集图像路径
    nc: 1 # 类别数，假设是一个类别
    names: [‘drone’] # 类别名称，假设是无人机

  • 注：train 和 val 路径应指向包含图像文件的文件夹。nc 是类别数，names 列表包含所有类别的名称（例如：’drone’ 表示无人机）

三.启动训练

  • python train.py –img 640 –batch 16 –epochs 50 –data /yolov5/data.yaml –weights yolov5s.pt –device cpu

  • 注：

    img 640：图像大小设置为 640x640。您可以根据需要调整图像大小。
    batch 16：每批次处理 16 张图片。您可以根据虚拟机的内存大小调整该值。CPU 训练时批量大小可以适当降低。
    epochs 50：训练 50 个轮次。您可以根据训练情况调整轮次。
    data ./yolov5_data/data.yaml：指定刚才创建的 data.yaml 配置文件。
    weights yolov5s.pt：使用预训练模型 yolov5s.pt 进行微调。yolov5s.pt 是 YOLOv5 的最小模型。
    device cpu：指定使用 CPU 进行训练。如果您使用 GPU，可以将其设置为 –device 0。

  • 之后命令行中看到训练的实时进度，包括损失函数、精度等指标。训练过程中，YOLOv5 会自动保存最佳模型（即在验证集上表现最好的模型）

运行验证

  • python val.py –weights runs/train/exp/weights/best.pt –data ./yolov5/data/data.yaml –img 640

四.模型推理

  • 训练完成后，可以使用训练好的模型进行无人机检测的推理（预测）：

    python detect.py –weights runs/train/exp/weights/best.pt –source ./yolov5/data/images/val –img 640

  • 注：

    weights：指定训练后的模型权重路径。
    source：指定待检测图像或视频的路径，您可以使用图像文件夹、单张图像或者视频文件。
    img 640：指定输入图像的大小。

  • 检测结果会保存在 runs/detect/exp 文件夹中。

模型选择

yolov5s.pt 是最小的模型，可以根据需要选择更大的模型（如 yolov5m.pt 或 yolov5l.pt）以获得更好的精度，但大模型训练速度较慢

五.模型部署

作者：zhangjn  创建时间：2025-04-13 14:22
最后编辑：zhangjn  更新时间：2025-04-16 09:02