深度学习检测算法YOLOv5的实战应用

news2024/11/20 3:36:19

在当前的检测项目中，需要一个高效且准确的算法来处理大量的图像数据。经过一番研究和比较，初步选择了YOLOv5作为算法工具。YOLOv5是一个基于深度学习的检测算法，以其快速和准确而闻名。它不仅能够快速处理图像数据，还能提供较高的检测准确率。

项目地址： https://github.com/ultralytics/yolov5
项目文档：https://docs.ultralytics.com/yolov5/tutorials/

配置环境

环境准备： 在进行YOLOv5的训练和预测之前，我们需要确保环境已经准备好。以下是具体的步骤：

创建虚拟环境：使用conda创建一个新的虚拟环境，命名为yolov5。
```
conda create -n yolov5 python=3.8
```
激活虚拟环境：激活刚刚创建的yolov5虚拟环境。
```
conda activate yolov5
```
安装依赖包：安装必要的依赖包，这里以GPU环境为例。
```
pip install ultralytics
```
克隆YOLOv5项目：从GitHub克隆YOLOv5项目。
```
git clone https://github.com/ultralytics/yolov5
```
进入项目目录：进入克隆的YOLOv5项目目录。
```
cd yolov5
```
安装项目依赖：安装项目中的依赖包。
```
pip install -r requirements.txt
```
注意：requirements.txt中的troch版本需要与你的硬件资源和CUDA版本相匹配。如果不匹配，可能无法启动GPU进行训练。

通过以上步骤，我们确保了环境已经准备好，可以顺利进行YOLOv5的训练和预测。接下来，将介绍如何进行训练。

训练

使用官方数据或者自己按照coco格式进行标注的数据都可以很简单的开始进行训练。如果是自我标注数据的话，数据格式：

images：文件夹里放原始图片；
labels：文件夹里放标注的标签文件；

训练是深度学习模型的核心环节，它决定了模型的性能和准确性。YOLOv5提供了强大的训练功能，允许用户根据自己的需求进行定制。以下是一个基本的训练命令示例，以及每个参数的含义：

python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml  --batch-size 128

-–data coco.yaml：指定训练数据集的配置文件。
-–epochs 300：设置训练的轮数（epoch）。
-–weights：指定预训练模型的权重，如果使用空字符串，则不加载预训练权重。
-–cfg yolov5n.yaml：指定配置文件，用于定义模型的结构和训练参数。
-–batch-size 128：设置每次迭代处理的图像数量

以上命令可以再现 YOLOv5 COCO 的效果。模型和数据集将从 YOLOv5 的最新版本中自动下载。在 V100 GPU 上，模型可选YOLOv5n/s/m/l/x，对应的训练时间分别为 1/2/4/6/8 天（使用多 GPU 训练将更快）。尽可能使用大的 --batch-size 值（16，32，64，128），或者传递 --batch-size -1 以启用 YOLOv5 的自动批量处理功能。显示的批量大小适用于 V100-16GB GPU。训练的时候也可以加入预训练模型和多GPU--weights ./pre-models/yolov5m.pt --device 0,1

可选的预训练模型如下表所示，按照精度和推理时间选择适合自己的模型：

所有检查点均按照默认设置训练了 300 个epoch。
Nano 和 Small 模型采用了 hyp.scratch-low.yaml 的超参数配置，而其他模型则采用了 hyp.scratch-high.yaml。
mAPval 值指的是在 COCO val2017 数据集上，单模型单尺度的评估结果。通过以下命令可以复现这一结果：python val.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65
速度是在 AWS p3.2xlarge 实例上，对 COCO val 图像集进行平均测量的。NMS 时间（约每张图像 1 毫秒）未计入其中。使用以下命令复现速度测试：python val.py --data coco.yaml --img 640 --task speed --batch 1
TTA（测试时间增强）包括反射和尺度增强。通过以下命令可以复现 TTA：python val.py --data coco.yaml --img 1536 --iou 0.7 --augment

预测

预测是深度学习模型的最终目的，它将训练好的模型应用于实际数据，以实现目标检测。YOLOv5提供了便捷的预测功能，可以快速对图像或视频进行目标检测。

使用项目已有的detect.py文件进行预测

detect.py可在各种不同的来源上执行推理任务，它会自动从最新的 YOLOv5 版本中下载所需的模型或者使用已训练好的模型，并将推断结果保存到 runs/detect 文件夹中。

python detect.py --weights yolov5s.pt --source  img.jpg   # image

weights：指定模型的权重文件，这里使用预训练的yolov5s模型，如果本地有的话就不下载，如果没有的话就网上进行下载；
source：除了图片，YOLOv5还支持多种输入源，包括视频、网络摄像头、屏幕截图等。可以根据需要选择合适的输入源。
- 0： # webcam
- img.jpg # image
- vid.mp4 # video
- screen # screenshot
- path/ # directory
- list.txt # list of images
- list.streams # list of streams
- 'path/*.jpg' # glob
- 'https://youtu.be/LNwODJXcvt4' # YouTube
- 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream

使用torch.hub进行预测

YOLOv5 PyTorch Hub 提供自动推理服务。所需的模型将从最新的 YOLOv5 版本自动下载。

import torch

# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")  # or yolov5n - yolov5x6, custom

# Images
img = "https://ultralytics.com/images/zidane.jpg"  # or file, Path, PIL, OpenCV, numpy, list

# Inference
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

上述内容就是从网络上下载模型并进行推理，如果想使用自己的模型的话进行推理，这里对其进行了封装，代码如下：

import os
import sys
import json
import torch
import numpy as np


CURRENT_DIR = os.path.abspath(os.path.dirname(__file__)) + '/'


class Detector:
    """implement detector
    """
    def __init__(self):
        # load model
        self.model = torch.hub.load(
            os.path.join(CURRENT_DIR, './'), 'custom',
            path=os.path.join(CURRENT_DIR, './yolov5s.pt'),
            source='local', device='cpu')

    def detect_img(self, img_file):
        """
        detect from input

        Args:
            file, Path, PIL, OpenCV, numpy, list
        """
        # inference
        results = self.model(img_file)
        crops = results.crop(save=False)  # cropped detections dictionary
        return crops


if __name__ == '__main__':
    dt = Detector()
    img = sys.argv[1]
    detect_res = dt.detect_img(img)
    print(detect_res)