Pytorch Yolov5训练自己目标检测模型和检测(CPU和GPU)

先来看看我们要实现的效果，我们将会通过数据来训练一个口罩检测的模型，并用pyqt5进行封装，实现图片口罩检测、视频口罩检测和摄像头实时口罩检测的功能。

代码下载代码的下载地址是：GitHub：GitHub - ultralytics/yolov5: YOLOv5 in PyTorch >

1、配置环境

配置anaconda虚拟环境

anaconda安装完成之后请切换到国内的源来提高下载速度 ，命令如下：

1
2
3
4
5
6

conda config --remove-key channels
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.bfsu.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes
pip config set global.index-url https://mirrors.ustc.edu.cn/pypi/web/simple

首先创建python3.8的虚拟环境，请在命令行中执行下列操作：

1
2

conda create -n pyY python==3.8.5
conda activate pyY

2、Pytorch:CPU和GPU的安装

实际测试情况是YOLOv5在CPU和GPU的情况下均可使用，不过在CPU的条件下训练那个速度会令人发指，所以有条件的小伙伴一定要安装GPU版本的Pytorch，没有条件的小伙伴最好是租服务器来使用。 在虚拟环境下安装pytorch版本地址为：Previous PyTorch Versions | PyTorch

1
2
3

conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cpuonly -c pytorch  # cpu
 
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge  # GPU

需要注意以下几点：

• 安装之前一定要先更新你的显卡驱动，去官网下载对应型号的驱动安装
• 30系显卡只能使用cuda11的版本
• 一定要创建虚拟环境，这样的话各个深度学习框架之间不发生冲突

安装完毕之后，我们来测试一下GPU是否

1
2
3

import torch
print(torch.__version__)
print(torch.cuda.is_available())

3、测试

1

在yolov5目录下执行下列代码，执行完毕之后会输出相关的信息，在runs目录下可以找到检测之后的结果

python detect.py –source data/images/bus.jpg –weights pretrained/yolov5s.pt

按照官方给出的指令，这里的检测代码功能十分强大，是支持对多种图像和视频流进行检测的，具体的使用方法如下：

1
2
3
4
5
6
7

python detect.py --source 0  # webcam
                           file.jpg  # image 
                           file.mp4  # video
                           path/  # directory
                           path/*.jpg  # glob
                           'https://youtu.be/NUsoVlDFqZg'  # YouTube video
                           'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream

4、数据处理

这里改成yolo的标注形式，之后专门出一期数据转换的内容。数据标注这里推荐的软件是labelimg，通过pip指令即可安装，在你的虚拟环境下执行

pip install labelimg 或

pip install labelimg -i https://mirror.baidu.com/pypi/simple

命令进行安装，然后在命令行中直接执行labelimg软件即可启动数据标注软件。然后数据标注（过程忽略大家自行搜索）注意：标志完成之后得到一系列的txt文件。这里的txt文件就是目标检测的标注文件，其中txt文件和图片文件的名称是一一对应的，打开具体的标注文件，你将会看到下面的内容，txt文件中每一行表示一个目标，以空格进行区分，分别表示目标的类别id，归一化处理之后的中心点x坐标、y坐标、目标框的w和h。

修改数据集配置文件：

标记完成的数据请按照下面的格式进行放置，方便程序进行索引。

YOLO_Mask

└─ data

├─ images

│ ├─ test # 下面放测试集图片

│ ├─ train # 下面放训练集图片

│ └─ val # 下面放验证集图片

└─ labels

├─ test # 下面放测试集标签

├─ train # 下面放训练集标签

├─ val # 下面放验证集标签

这里的配置文件是为了方便我们后期训练使用，我们需要在data目录下创建一个mask_data.yaml的文件，如下图所示：

# Custom data for safety helmet

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]

train: C:/score/images/train # 训练集位置

val: C:/score/images/val # 验证机位置

# number of classes

nc: 2 # 数据集类别

# class names

names: [‘fire’, ‘smoke’] # 目标类名称

模型训练

模型基本训练

在models下建立一个mask_yolov5s.yaml（复制yolov5s.yaml改为名字）的模型配置文件。整理好数据集配置文件、模型配置文件、预训练模型执行下列代码运行程序即可：

1
2
3

python train.py --data mask_data.yaml --cfg mask_yolov5s.yaml --weights pretrained/yolov5s.pt --epoch 100 --batch-size 4 --device cpu
 
python train.py --data mask_data.yaml --cfg mask_yolov5s.yaml --weights pretrained/yolov5s.pt --epoch 100 --batch-size 16 --device 0

训练代码成功执行之后在命令行中输出下列信息，接下来就是安心等待模型训练结束。

在train/runs/exp3的目录下可以找到训练得到的模型和日志文件，当然还有一些各种指标的图片，大家自从查看。

模型评估

其中目标检测最常用的评价指标是mAP，mAP是介于0到1之间的一个数字，这个数字越接近于1，就表示你的模型的性能更好。一般我们会接触到两个指标，分别是召回率recall和精度precision，两个指标p和r都是简单地从一个角度来判断模型的好坏，均是介于0到1之间的数值，其中接近于1表示模型的性能越好，接近于0表示模型的性能越差，为了综合评价目标检测的性能，一般采用均值平均密度map来进一步评估模型的好坏。我们通过设定不同的置信度的阈值，可以得到在模型在不同的阈值下所计算出的p值和r值，一般情况下，p值和r值是负相关的，绘制出来可以得到如下图所示的曲线，其中曲线的面积我们称AP，目标检测模型中每种目标可计算出一个AP值，对所有的AP值求平均则可以得到模型的mAP值，以本文为例，我们可以计算火和烟雾的两个目标的AP值，我们对两组AP值求平均，可以得到整个模型的mAP值，该值越接近1表示模型的性能越好。

如果你的目录下没有这样的曲线，可能是因为你的模型训练一半就停止了，没有执行验证的过程，你可以通过下面的命令来生成这些图片。

python val.py –data data/mask_data.yaml –weights runs/train/exp_yolov5s/weights/best.pt –img 640

模型使用：

模型的使用全部集成在了detect.py目录下，你按照下面的指令指你要检测的内容即可

1
2
3
4
5
6
7
8
9
10
11
12

# 检测摄像头
python detect.py  --weights runs/train/exp_yolov5s/weights/best.pt --source 0  # webcam
# 检测图片文件
 python detect.py  --weights runs/train/exp_yolov5s/weights/best.pt --source file.jpg  # image 
# 检测视频文件
  python detect.py --weights runs/train/exp_yolov5s/weights/best.pt --source file.mp4  # video
# 检测一个目录下的文件
 python detect.py --weights runs/train/exp_yolov5s/weights/best.pt path/  # directory
# 检测网络视频
 python detect.py --weights runs/train/exp_yolov5s/weights/best.pt 'https://youtu.be/NUsoVlDFqZg'  # YouTube video
# 检测流媒体
 python detect.py --weights runs/train/exp_yolov5s/weights/best.pt 'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream                  

比如以我们的模型为例，如果我们执行python detect.py –weights runs/train/exp_yolov5s/weights/best.pt –source data/images/fishman.jpg的命令便可以得到这样的一张检测结果。

转载自知乎：Pytorch Yolov5训练自己目标检测模型和检测(CPU和GPU) - 知乎 (zhihu.com)