《计算机视觉之PyTorch数字图像处理》--目标检测之YOLOv5预训练模型、YOLOv5模型训练

dirty

《计算机视觉之PyTorch数字图像处理》--目标检测之YOLOv5预训练模型、YOLOv5模型训练 [复制链接]

      前面讲了Torchvision中的预训练模型，并对图像进行了检测。本篇讲述YOLOv5预训练模型、YOLOv5模型及其训练。

 

YOLOv5预训练模型

      YOLO系列具有检测速度快、精度也不低的优点，已经成为了最常用的目标检测模型。目前主流的YOLO模型是

链接已隐藏，如需查看请登录或者注册

。YOLOv5安装如下三步骤：

(1)Git克隆一份到本地git clone 

链接已隐藏，如需查看请登录或者注册

 ,资源如下：

(2)下载权重文件。网址:

链接已隐藏，如需查看请登录或者注册

. 在权重文件下，找到Assets目录，在该目录下即列出了YOLOv5系列所有36个模型的权重参数。

(3)安装项目所需的第三方库。 这里需要而外安装ipython、mss、albumentations、pycocotools库，修改requirements.txt如下。命令安装库 pip install -r .\requirements.txt 

     

      安装好后，使用YOLOv5模型非常简单，只需要运行detect.py设置需要检测的数据源（source)即可。

YOLOv5支持目标的数据源有：

●Webcam（连接到计算机的摄像头） ：(OPTION = 0) 实时的从摄像头读取图像并进行目标检测；

●图像： (OPTION = filename.jpg) 对指定路径中的图像进行检测，并输出检测结果到新的图像文件；

●视频： (OPTION = filename.mp4) 对指定路径下的视频进行检测，创建一个带有目标检测结果的视频；

●目录： (OPTION = directory_name/) 对一个目录下的所有的指定文件检测，并保存带有检测结果；

●RTSP流：(OPTION = rtsp://170.93.143.139/rtplive/470011e600ef003696235daa ) 从RTSP流中进行目标检测；

●RTMP流： (OPTION = rtmp://192.168.1.105/live/test) 从RTMP流中进行目标检测；

●HTTP流：(OPTION = http://112.50.243.8/PLTV/88888/224/3221225900/1.m3u8) 从HTTP流中进行目标检测。

 

      下面以一张图片进行YOLOv5预训练模型检测。操作与生成如下：

      检测结果如下，上图原图，下图检测结果

 

      我们再用前面Torchvision预训练模型检测的图片在此检测，可以看到YOLOv5检测到了更多的目标。

 

YOLOv5模型及其训练

数据集的构建

      YOLOv5将数据集的创建包含到项目本身，不需要使用者自定义数据集，只需要按照YOLOv5对数据集格式的要求整理好数据和标签，并使用配置文件记录数据集的位置和类别即可。数据集的配置文件存放于yolov5文件夹下的data子文件夹中，并且以.yaml的格式进行存储。YAML文件是一种用于记录配置的标记语言，YOLOv5使用该格式存储数据集和模型的配置。YAML文件路径及配置文件按内容参考如下。path参数指明了该数据集的位置；train参数指向了包含所有图像训练样本路径的文本文件；val和test参数的含义与train参数含义相同，只是对验证集和测试集的区别；names指明了数据信中包含的所有类别的标签和对应的释义。

      对于样本的标签有多种存储方式，这里我们介绍其中一种。YOLOv5可以将样本图像的标签信息存储到与图像名称相同，但结尾为.txt的文本中，将标签文件与图像文件都放置在同一个目录train文件夹下。YOLOv5定义的一种目标检测标签文件格式，在该文件中图像中的每个目标使用一行记录，包含用空格分隔的5个数值，分别是目标的类别标签值、归一化的目标中心点x坐标、归一化的目标中心点y坐标、归一化的目标的宽度和归一化的目标的高度。

 

模型的构建

      和数据集定义方式相似，检测模型的构建也使用YAML配置文件的方式进行。YOLOv5模型的定义保存在models文件夹下，按照不同的规模每个模型使用一个yolov5*.yaml进行定义。这里选择yolov5s.yaml定义的YOLOv5s模型展开讲述自定义模型方法：

      首先复制yolov5s.yaml文件，修改文件名为yolov5s_copy.yaml；

      其次，打开yolov5slslm.yaml文件，只需要将表示类别数的nc参数修改为2后，保存退出。

      了解模型参数说明：类别数的nc；depth_multiple参数定义了模型的深度系数（模型包含的层数）；width_multiple参数定义了模型中各层通道的数量；anchors参数定义了在三个不同尺度特征图像上的锚框，每个尺度上有三组不同形状的锚框，以表示三个框的宽和高；backbone定义了以特征提取为主的骨干网络的结构，head定义了检测模型三个尺度特征的融合和检测输出。

      YOLOv5模型的定义使用了配置文件的方式，在底层YOLOv5会解析并转化为Pytroch的模型。对于模型创建进一步的了解可以查看models目录下的yolo.py和common.py文件中的源代码。

      YOLOv5对于模型的训练已经集成到项目代码中，在准备好数据集和模型后，直接调用train.py进行训练即可。设置训练轮数为20轮，批大小为4 的命令为：

python .\train.py --cfg yolov5slslm.yaml  --data lslm.yaml --epochs 20 --batch-size 4

      执行上述命令，就开始了对螺丝螺母数据集的训练，在训练过程中终端里会显示训练的全部参数设置，模型的结构，训练结果的保存目录，以及详细的训练过程。程序会将结果保存到runs目录下的train文件夹中的一个名为exp*的文件夹。

      在结果文件夹中，weights文件夹中保存了模型最后一轮训练的权重文件last.pt和模型训练过程中最优的权重文件best.pt。当需要进行检测图像中的螺丝和螺母时，就需要使用训练得到的权重文件。训练完成后，进行测试就十分简单，直接使用YOLOv5预训练模型的方法，调用detect.py进行即可：

python .\detect.py --weights ./runs/train/exp1/weights/best.pt --source D:\data\lslm\test，检测结果会保存在runs目录下的detect文件夹中的名为exp*的目录中。打开该目录就可以看到检测的结果。

 

       通过本篇学习实践使用YOLOv5预训练模型检测，以及对YOLOv5数据集、模型构建方法过程有了一个清晰梳理，对工程实践很有助益。