利用yolov5实现口罩佩戴检测算法

Python之王

共 16353字,需浏览 33分钟

 ·

2022-02-23 12:21


首先申明,这篇博客是用于记录我第一次完全从头到尾跑通一个算法,我会在此博客详细写出我的具体过程,以供大家参考,可能会和炮哥博客有些重合,没办法毕竟我就是用他的博客来训练模型的。但这篇博客我会结合炮哥的博客和我自己训练过程中的一些问题和心得来写,所以还是会有所不的!!!博主其实也是个深度学习的小菜鸟hhh。在此感谢。

一、必要的环境依赖和项目克隆

1.1环境的安装

这块主要是深度学习的环境安装,说实话我觉得这块也挺麻烦,我有时候一弄就是两三天,总会出现奇奇怪怪的问题,在这我推荐一篇,个人觉得已经写的非常详细了。不过里面的paddlepaddle我并没有安装,我没有用到,大家可以自己选择是否安装。这里我重新创建了一个yolov5的环境用来跑算法。 这里博主感叹下,不同环境可以装不同版本cuda,之前博主反反复复在自己系统安装不同版本的cuda,搞了我三四天都不行,真的晕了。最后发现原来可以这样,这么方便!!!

1.2项目的克隆

首先去github的官网(这是国外的网站,有时候进的去有时候进不去,多试几次)界面如下 点击code,将整个代码下载下来

1.3项目整体介绍

将下载的压缩包解压,用pycharm打开 注:这块我是参考炮哥的,毕竟我对yolov5的结构和原理并不是非常清楚,只是有所了解。 现在来对代码的整体目录做一个介绍:

├── data:主要是存放一些超参数的配置文件(这些文件(yaml文件)是用来配置训练集和测试集还有验证集的路径的,其中还包括目标检测的种类数和种类的名称);还有一些官方提供测试的图片。如果是训练自己的数据集的话,那么就需要修改其中的yaml文件。但是自己的数据集不建议放在这个路径下面,而是建议把数据集放到yolov5项目的同级目录下面。

├── models:里面主要是一些网络构建的配置文件和函数,其中包含了该项目的四个不同的版本,分别为是s、m、l、x。从名字就可以看出,这几个版本的大小。他们的检测测度分别都是从快到慢,但是精确度分别是从低到高。这就是所谓的鱼和熊掌不可兼得。如果训练自己的数据集的话,就需要修改这里面相对应的yaml文件来训练自己模型。

├── utils:存放的是工具类的函数,里面有loss函数,metrics函数,plots函数等等。

├── weights:放置训练好的权重参数。

├── detect.py:利用训练好的权重参数进行目标检测,可以进行图像、视频和摄像头的检测。

├── train.py:训练自己的数据集的函数。

├── test.py:测试训练的结果的函数。

├──requirements.txt:这是一个文本文件,里面写着使用yolov5项目的环境依赖包的一些版本,可以利用该文本导入相应版本的包。

以上就是yolov5项目代码的整体介绍。我们训练和测试自己的数据集基本就是利用到如上的代码。

1.4依赖包的下载

打开requirements.txt 可以看到里面有很多需要包的版本,而这你只需在终端运行pip install -r requirements.txt命令即可 而这大概大多包都没问题,有些人可能会在装pycocotools>=2.0这个包有问题,这里博主采用的是安装visual studio C++ build tools (这好像也能解决Microsoft Visual C++ 14.0 is required的报错问题),再在(项目的环境下,我这里是yolov5)终端pip install pycocotools 就安装成功了。附上一位大好人给的安装包,多亏了他让我把这个包拿下!链接:https://pan.baidu.com/s/1GKjbxs9_y-Z4nkGGWsKqPA 提取码:1ors 安装很简单,解压后是一个iso文件双击,进去后点exe安装就ok了 注:这是我的方法,我不保证大家都能利用这个解决

到这,深度学习所需的环境和依赖包就准备好了。

二、数据集和预训练权重的准备

2.1数据集的准备

如果你有图片准备自己打标签的话可以利用labelimg来制作自己的数据集,这里推荐炮哥又一篇博客,

在这附上博主用的口罩数据集链接:https://pan.baidu.com/s/1Gud8jemSCdjG00TYA74WpQ 提取码:sv74 这里的lables是已经是txt(yolo的训练标签就是txt),而一般的标签都是xml格式。标签:0:no-mask,1:mask 注:这里大概8000张图片,包含戴口罩和不戴口罩的,我从里面抽了2000张,不抽多的原因是博主的显卡不行(GTX1050),所以就只用了1/4的数据。

2.2将数据集划分为训练集和验证集

这呢,推荐大家去看炮哥的这篇博客, 这里因为博主用的数据集因为标签已经是txt格式了,但我先将txt转xml格式,再用代码直接将xml格式转为yolo(txt)格式并划分训练集和测试集。(这里不直接用txt的格式直接划分,炮哥是这样解释的,txt划分后放入训练会出错) 简单的说就是先用炮哥博客里的2代码将txt转为xml,如下图: 再利用1的代码 其实炮哥的博客讲的已经非常详细了,只要格式放的没问题就能得到这样的划分结果如下图所示: 其中多的两个文件(train.cache,val.cache)是因为博主已经训练完毕多的。划分结束后将VOCdevkit整个文件夹放到yolov5的代码中 注意: 这里如果在转换yolo格式(txt)到VOC格式(xml)时出现 KeyError: 'None’的报错,但一部分的txt文件已经转换成xml,这时候只需将下一个待转换图片和txt(对应的标签)删除即可。(估计是文件有问题,删除后就可以将后面的全部转换了)

2.3获得预训练权重

为什么要用预训练权重呢,一般是为了缩短网络的训练时间,达到更好的精度。而yolov5的5.0版本给我们提供了几个预训练权重,我们可以对应我们不同的需求选择不同的版本的预训练权重。通过如下的图可以获得权重的名字和大小信息,可以预料的到,预训练权重越大,训练出来的精度就会相对来说越高,但是其检测的速度就会越慢。预训练权重可以通过github进行下载,,点进去后往下拉就能找到。本次训练自己的口罩数据集用的预训练权重为yolov5s.pt。

将下载的权重放到文件夹weights下 到此数据集和权重就已经准备好了接下来就可以准备开始训练自己的yolov5口罩检测模型了。

三、训练口罩检测模型

3.1相关文件的配置

在开始训练前还需要对项目里的相关文件进行修改,一个是数据配置文件,另一个是模型配置文件 首先是数据配置文件,在data下找到voc.yaml,将其复制一份再重命名为mask.yaml 将mask.yaml文件里的4个箭头出进行修改,对第一个箭头指向的代码加注释(这里已经注释),按炮哥说的,不注释的话训练时候会报错。箭头2中需要将之前划分好的训练集和测试集的路径填上。第三个箭头填写需要检测的类别数,因为数据集的标签是no-mask和mask。所以这里填两类。第四个箭头中填写类别的名称,就用标签的名称(no-mask和mask)。这样这个yaml文件就改好了。 接下来改模型的配置文件,因为该项目使用的yolov5s.pt这个权重,所以使用models下的yolov5s.yaml(不同的预训练权重对应不同的网络层数,用错会报错)。同理复制yolov5.yaml为mask.yaml,打开文件进行参数修改 这里只要修改识别的类别数即可,为两类 到这里文件的配置就已经好了

3.2准备训练模型

找到train.py文件 找到main函数入口,修改几个参数,模型的主要参数解析参考炮哥。

if __name__ == '__main__':
"""
    opt模型主要参数解析:
    --weights:初始化的权重文件的路径地址
    --cfg:模型yaml文件的路径地址
    --data:数据yaml文件的路径地址
    --hyp:超参数文件路径地址
    --epochs:训练轮次
    --batch-size:喂入批次文件的多少
    --img-size:输入图片尺寸
    --rect:是否采用矩形训练,默认False
    --resume:接着打断训练上次的结果接着训练
    --nosave:不保存模型,默认False
    --notest:不进行test,默认False
    --noautoanchor:不自动调整anchor,默认False
    --evolve:是否进行超参数进化,默认False
    --bucket:谷歌云盘bucket,一般不会用到
    --cache-images:是否提前缓存图片到内存,以加快训练速度,默认False
    --image-weights:使用加权图像选择进行训练
    --device:训练的设备,cpu;0(表示一个gpu设备cuda:0);0,1,2,3(多个gpu设备)
    --multi-scale:是否进行多尺度训练,默认False
    --single-cls:数据集是否只有一个类别,默认False
    --adam:是否使用adam优化器
    --sync-bn:是否使用跨卡同步BN,在DDP模式使用
    --local_rank:DDP参数,请勿修改
    --workers:最大工作核心数
    --project:训练模型的保存位置
    --name:模型保存的目录名称
    --exist-ok:模型目录是否存在,不存在就创建
"
""
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights'type=str, default='yolov5s.pt'help='initial weights path')
    parser.add_argument('--cfg'type=str, default=''help='model.yaml path')
    parser.add_argument('--data'type=str, default='data/coco128.yaml'help='data.yaml path')
    parser.add_argument('--hyp'type=str, default='data/hyp.scratch.yaml'help='hyperparameters path')
    parser.add_argument('--epochs'type=int, default=300)
    parser.add_argument('--batch-size'type=int, default=16, help='total batch size for all GPUs')
    parser.add_argument('--img-size', nargs='+'type=int, default=[640, 640], help='[train, test] image sizes')
    parser.add_argument('--rect', action='store_true'help='rectangular training')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true'help='only save final checkpoint')
    parser.add_argument('--notest', action='store_true'help='only test final epoch')
    parser.add_argument('--noautoanchor', action='store_true'help='disable autoanchor check')
    parser.add_argument('--evolve', action='store_true'help='evolve hyperparameters')
    parser.add_argument('--bucket'type=str, default=''help='gsutil bucket')
    parser.add_argument('--cache-images', action='store_true'help='cache images for faster training')
    parser.add_argument('--image-weights', action='store_true'help='use weighted image selection for training')
    parser.add_argument('--device', default=''help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true'help='vary img-size +/- 50%%')
    parser.add_argument('--single-cls', action='store_true'help='train multi-class data as single-class')
    parser.add_argument('--adam', action='store_true'help='use torch.optim.Adam() optimizer')
    parser.add_argument('--sync-bn', action='store_true'help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--local_rank'type=int, default=-1, help='DDP parameter, do not modify')
    parser.add_argument('--workers'type=int, default=8, help='maximum number of dataloader workers')
    parser.add_argument('--project', default='runs/train'help='save to project/name')
    parser.add_argument('--entity', default=None, help='W&B entity')
    parser.add_argument('--name', default='exp'help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true'help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true'help='quad dataloader')
    parser.add_argument('--linear-lr', action='store_true'help='linear LR')
    parser.add_argument('--label-smoothing'type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--upload_dataset', action='store_true'help='Upload dataset as W&B artifact table')
    parser.add_argument('--bbox_interval'type=int, default=-1, help='Set bounding-box image logging interval for W&B')
    parser.add_argument('--save_period'type=int, default=-1, help='Log model after every "save_period" epoch')
    parser.add_argument('--artifact_alias'type=str, default="latest"help='version of dataset artifact to be used')
    opt = parser.parse_args()

将yolov5s.pt的权重的相对路径放在这 这里路径只需选中文件右击,再点击复制路径即可 点击图中蓝色复制得到相对路径 这里将之前改的models下的mask.yaml路径放入 这里放data下的mask.yaml 这两个参数大家根据自己电脑的配置修改,第一个epochs是训练次数,我这只训练200次,第二个batch-size是每次输入图片的数量,我这只能选4,多了就会报CUDA out of memory的错误。 还有就是CPU的工作核心,我的cpu是4核所以这里改为4 上面都设置好后就可以训练了,但根据炮哥博客,pycharm的用户会出现下面的报错。这里给上炮哥的解决方法 在utils路径下找到datasets.py这个文件,将里面的81行里面的参数num_workers改成0 这里运行train.py文件估计还会报一个错 这里少了SPPF类,需要到github下载yolov5-6.0,打开文件找到models文件下的common.py,到里面复制SPPF类,并将这段代码(这里我已经放在下面了)复制到自己项目文件的models/common.py里去,具体位置应该没要求,我放在了149的SPP类之后。

class SPPF(nn.Module):
    # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
    def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
            y1 = self.m(x)
            y2 = self.m(y1)
            return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))

到这里,就可以运行train.py文件开始训练自己的模型了。博主当时改到这,生怕报别的错,最后开始训练了,深深呼出了一口气!开始训练后是这样的 训练结束,2000张图片,200轮共耗时12h!

3.3启用tensorbord查看训练结果

训练结束后在终端输入如下命令

tensorboard --logdir=runs

复制网址在浏览器打开就可以看到了,可以看到博主的口罩检测模型训练结果还行

四、推理测试

训练结束后,会产生一个runs的文件夹,在runs/train/exp3/weights会产生两个权重文件,这里exp3的原因是因为博主前面两次训练失败了,如果大家一次就训练成功应该是exp的文件夹下。其中best.pt(最好的权重),last.pt(是最后一轮的权重),我们推理利用最好的权重(best.pt)。 找到目录下的detect.py文件并打开 同样的找到main函数,这里有模型的主要参数。模型的主要参数解析同样参考的炮哥hh

f __name__ == '__main__':
"""
--weights:权重的路径地址
--source:测试数据,可以是图片/视频路径,也可以是'0'(电脑自带摄像头),也可以是rtsp等视频流
--output:网络预测之后的图片/视频的保存路径
--img-size:网络输入图片大小
--conf-thres:置信度阈值
--iou-thres:做nms的iou阈值
--device:是用GPU还是CPU做推理
--view-img:是否展示预测之后的图片/视频,默认False
--save-txt:是否将预测的框坐标以txt文件形式保存,默认False
--classes:设置只保留某一部分类别,形如0或者0 2 3
--agnostic-nms:进行nms是否也去除不同类别之间的框,默认False
--augment:推理的时候进行多尺度,翻转等操作(TTA)推理
--update:如果为True,则对所有模型进行strip_optimizer操作,去除pt文件中的优化器等信息,默认为False
--project:推理的结果保存在runs/detect目录下
--name:结果保存的文件夹名称
"
""
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+'type=str, default='yolov5s.pt'help='model.pt path(s)')
    parser.add_argument('--source'type=str, default='data/images'help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size'type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres'type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres'type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default=''help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true'help='display results')
    parser.add_argument('--save-txt', action='store_true'help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true'help='save confidences in --save-txt labels')
    parser.add_argument('--nosave', action='store_true'help='do not save images/videos')
    parser.add_argument('--classes', nargs='+'type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true'help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true'help='augmented inference')
    parser.add_argument('--update', action='store_true'help='update all models')
    parser.add_argument('--project', default='runs/detect'help='save results to project/name')
    parser.add_argument('--name', default='exp'help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true'help='existing project/name ok, do not increment')
    opt = parser.parse_args()

在这里将best.pt(最好的权重)的路径输入

4.1对图片视频测试

在这里输入对所需测试的图片的路径,然后运行detect.py 就会在runs/detect/exp下显示结果 这是原图 这是结果图,可以看到效果还是不错的 这里是没戴口罩的,可以看到检测时间很快0.238s

对视频的测试就是将图片的路径换成视频的路径而已。

4.2用摄像头测试

要利用摄像头测试只需将路径改写成0即可。

这里根据炮哥的博客会报错还需要对utils下的datasets.py文件进行修改 找到第279行,将两个url加上str() 摄像头的效果图就是这样,摘下口罩会显示no-mask

最后

这篇博客大概写了4,5天吧,差不多把基本的细节都写了,说实话如果大家能掌握这一套流程,基本以后想要训练其他关于分类的模型,其实都差不多,像炮哥给的安全帽,或者其他一些更多的类别之类。都能上手操作。

链接:https://pan.baidu.com/s/1r0IoJLtMZIodQoRcIxmNwA 提取码:zuu2




Python“宝藏级”公众号【Python之王】专注于Python领域,会爬虫,数分,C++,tensorflow和Pytorch等等

近 2年共原创 100+ 篇技术文章。创作的精品文章系列有:

日常收集整理了一批不错的 Python 学习资料,有需要的小伙可以自行免费领取。

获取方式如下:公众号回复资料领取Python等系列笔记,项目,书籍,直接套上模板就可以用了。资料包含算法、python、算法小抄、力扣刷题手册和 C++ 等学习资料!

浏览 99
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报