YOLOv5深度学习检测算法实战指南

YOLOv5深度学习检测算法实战指南

YOLOv5是一个基于深度学习的目标检测算法，以其快速和准确而闻名。它不仅能够快速处理图像数据，还能提供较高的检测准确率。本文将详细介绍YOLOv5的环境配置、训练和预测过程，并提供具体的代码示例。

配置环境

在进行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中的torch版本需要与你的硬件资源和CUDA版本相匹配。如果不匹配，可能无法启动GPU进行训练。

训练

使用官方数据或者自己按照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

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

1713753606217.png
1713752328003.png

• 所有检查点均按照默认设置训练了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)

这里使用results.crop是因为这个返回的信息会比较多，返回的结果包含预测的标签，置信度，以及检测区域等，便于后续业务逻辑处理。

结论

YOLO系列算法是检测算法里面非常实用的一种工具项目，能够在工业界得到很好的使用。本文介绍了一下Yolov5的训练及预测流程，便于大家使用以及后续自己回顾使用。