基于yolo的口罩识别(开源代码和数据集)

2020年开头真的很人意外，开年爆发了疫情。此次疫情牵动了各行各业，在这里衷心的感谢奋斗在一线的医疗工作者：您们辛苦了。作为一名非医专业的学生，在这样情况下，除了不乱跑以外，我也想以另一种方式去致敬那些保护着全国人民安全的工作人员。接下来就来介绍本项目的开发过程。

1. 开发环境

笔者的开发环境如下：
ubuntu16.04
PyTorch 1.1.0
anaconda
opencv-python
tqdm
matplotlib
pycocotools

2. 收集数据集

本项目的数据集是笔者花了一个晚上通过网络爬虫的方式收集的（ps:没有违反爬虫协定）。笔者通过使用python写了一个小爬虫，爬虫的代码如下：

import re
import requests
from urllib import error
from bs4 import BeautifulSoup
import os

num = 0
numPicture = 0
file = ''
List = []


def Find(url):
    global List
    print('正在检测图片总数，请稍等.....')
    t = 0
    i = 1
    s = 0
    while t < 1000:
        Url = url + str(t)
        try:
            Result = requests.get(Url, timeout=7)
        except BaseException:
            t = t + 60
            continue
        else:
            result = Result.text
            pic_url = re.findall('"objURL":"(.*?)",', result, re.S)  # 先利用正则表达式找到图片url
            s += len(pic_url)
            if len(pic_url) == 0:
                break
            else:
                List.append(pic_url)
                t = t + 60
    return s


def recommend(url):
    Re = []
    try:
        html = requests.get(url)
    except error.HTTPError as e:
        return
    else:
        html.encoding = 'utf-8'
        bsObj = BeautifulSoup(html.text, 'html.parser')
        div = bsObj.find('div', id='topRS')
        if div is not None:
            listA = div.findAll('a')
            for i in listA:
                if i is not None:
                    Re.append(i.get_text())
        return Re


def dowmloadPicture(html, keyword):
    global num
    # t =0
    pic_url = re.findall('"objURL":"(.*?)",', html, re.S)  # 先利用正则表达式找到图片url
    print('找到关键词:' + keyword + '的图片，即将开始下载图片...')
    for each in pic_url:
        print('正在下载第' + str(num + 1) + '张图片，图片地址:' + str(each))
        try:
            if each is not None:
                pic = requests.get(each, timeout=7)
            else:
                continue
        except BaseException:
            print('错误，当前图片无法下载')
            continue
        else:
            string = file + r'\\' + keyword + '_' + str(num) + '.jpg'
            fp = open(string, 'wb')
            fp.write(pic.content)
            fp.close()
            num += 1
        if num >= numPicture:
            return


if __name__ == '__main__':  # 主函数入口
    tm = int(input('请输入每类图片的下载数量 '))
    numPicture = tm
    line_list = []
    with open('./name.txt', encoding='utf-8') as file:
        line_list = [k.strip() for k in file.readlines()]  # 用 strip()移除末尾的空格

    for word in line_list:
        url = 'http://image.baidu.com/search/flip?tn=baiduimage&ie=utf-8&word=' + word + '&pn='
        tot = Find(url)
        Recommend = recommend(url)  # 记录相关推荐
        print('经过检测%s类图片共有%d张' % (word, tot))
        file = word
        y = os.path.exists(file)
        if y == 1:
            print('该文件已存在，请重新输入')
            file = word
            os.mkdir(file)
        else:
            os.mkdir(file)
        t = 0
        tmp = url
        while t < numPicture:
            try:
                url = tmp + str(t)
                result = requests.get(url, timeout=10)
                print(url)
            except error.HTTPError as e:
                print('网络错误，请调整网络后重试')
                t = t + 60
            else:
                dowmloadPicture(result.text, word)
                t = t + 60
        numPicture = numPicture + tm

    print('任务完成')

运行此爬虫需要配以一个name.txt,在其中写下所需要爬取的信息即可。，rename.py是对获取的数据进行命名整理的。代码如下：

import os

path_name='./JPEGImages'
#path_name :表示你需要批量改的文件夹
i=0
for item in os.listdir(path_name):#进入到文件夹内，对每个文件进行循环遍历
    os.rename(os.path.join(path_name,item),os.path.join(path_name,('masks_'+'00'+str(i)+'.jpg')))#os.path.join(path_name,item)表示找到每个文件的绝对路径并进行拼接操作
    i+=1

文件结构如下：

3.数据的清洗

这里笔者采用的数据清洗就比较的笨了，就是人为一个一个的筛选，整理了一个晚上，请原谅笔者的愚蠢。

4. 标注数据集

标注数据集同数据的清洗而言就比较简单了，比较的机械化。这里笔者使用labelImg这个标注工具进行标注。这里做的是简单的二分类任务，即：Mask和unMasked.标注比数据清洗要快一点。

5. 相关准备

准备工作的话就是使用一个版本的大体框架，这里使用的是yolov3，yolov3框架。下载后可以重新命名，当然也可不命名，这个随性就可以了。

6. 数据装载

申明一下最后会将数据集提供给大家。
我们需要将数据集Annotations、JPEGImages复制到YOLOV3工程目录下的data文件下；同时新建两个文件夹，分别命名为ImageSets和labels，最后我们将JPEGImages文件夹复制粘贴一下，并将文件夹重命名为images

7. 代码的构建

这里需要在项目的根目录下新建连个python文件：makeTxt.py和voc_label.py 其中makeTxt.py是用于生成data/ImageSets下的test.txt、train.txt、trainval.txt以及val.txt的;而voc_label.py是将Annotations里数据集的标注文件转为data/label下txt格式的。
makeTxt.py的代码如下：

import os
import random
 
trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'data/Annotations'
txtsavepath = 'data/ImageSets'
total_xml = os.listdir(xmlfilepath)
 
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
 
ftrainval = open('data/ImageSets/trainval.txt', 'w')
ftest = open('data/ImageSets/test.txt', 'w')
ftrain = open('data/ImageSets/train.txt', 'w')
fval = open('data/ImageSets/val.txt', 'w')
 
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
 
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

voc_label.py的代码如下：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
 
sets = ['train', 'test','val']
 
classes = ["mask","unmask"]   #我们只是检测细胞，因此只有一个类别
 
 
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)
 
 
def convert_annotation(image_id):
    in_file = open('data/Annotations/%s.xml' % (image_id))
    out_file = open('data/labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
 
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
 
 
wd = getcwd()
print(wd)
for image_set in sets:
    if not os.path.exists('data/labels/'):
        os.makedirs('data/labels/')
    image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()
    list_file = open('data/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write('data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

分别运行makeTxt.py和voc_label.py会在data/ImageSets的四个文件中出现如下变化

再运行voc_label.py会在label文件夹下产生如下的变化。

接着还要配置两个文件
在data文件下新建rbc.data，配置内容如下：

classes=2
train=data/train.txt
valid=data/test.txt
names=data/rbc.names
backup=backup/
eval=coco

值得注意的是这里是对人是否佩戴口罩进行识别的，即分为两类佩戴口罩和为佩戴口罩。所以这里需要将classes设置为2.
再在data文件下新建rbc.names，配置内容如下：

Masking
unMasked

8. 修改配置文件

这里需要将cfg下的yolov3-tiny.cfg文件进行修改，修改内容如下：

[net]
# Testing
batch=1
subdivisions=1
# Training
# batch=64
# subdivisions=2
width=416
height=416
channels=3
momentum=0.9
decay=0.005
angle=0
saturation = 1.5
exposure = 1.5
hue=.1

learning_rate=0.005
burn_in=1000
max_batches = 500200
policy=steps
steps=400000,450000
scales=.1,.1

[convolutional]
batch_normalize=1
filters=16
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=2

[convolutional]
batch_normalize=1
filters=32
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=2

[convolutional]
batch_normalize=1
filters=64
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=2

[convolutional]
batch_normalize=1
filters=128
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=2

[convolutional]
batch_normalize=1
filters=256
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=2

[convolutional]
batch_normalize=1
filters=512
size=3
stride=1
pad=1
activation=leaky

[maxpool]
size=2
stride=1

[convolutional]
batch_normalize=1
filters=1024
size=3
stride=1
pad=1
activation=leaky

###########

[convolutional]
batch_normalize=1
filters=256
size=1
stride=1
pad=1
activation=leaky

[convolutional]
batch_normalize=1
filters=512
size=3
stride=1
pad=1
activation=leaky

[convolutional]
size=1
stride=1
pad=1
filters=21   ####因为这里是二分类，根据论文中的介绍这里需要根据3*(5+classnum),应此为21
activation=linear



[yolo]
mask = 3,4,5
anchors = 10,14,  23,27,  37,58,  81,82,  135,169,  344,319
classes=2    ### 这里因为是二分类，所以改为2
num=6
jitter=.3
ignore_thresh = .7
truth_thresh = 1
random=1

[route]
layers = -4

[convolutional]
batch_normalize=1
filters=128
size=1
stride=1
pad=1
activation=leaky

[upsample]
stride=2

[route]
layers = -1, 8

[convolutional]
batch_normalize=1
filters=256
size=3
stride=1
pad=1
activation=leaky

[convolutional]
size=1
stride=1
pad=1
filters=21   ### 因为这里是二分类，根据论文中的介绍这里需要根据3*(5+classnum),应此为21
activation=linear

[yolo]
mask = 0,1,2
anchors = 10,14,  23,27,  37,58,  81,82,  135,169,  344,319
classes=2    ### 这里因为是二分类，所以改为2
num=10
jitter=.3
ignore_thresh = .7
truth_thresh = 1
random=1

9. 训练数据集

完成了上述操作后，即可进行数据集的训练操作了。
在根目录下打开命令行，输入如下命令：

python train.py --data-cfg data/rbc.data --cfg cfg/yolov3-tiny.cfg --epochs 100

参数说明：
rbc.data 是data文件夹下
yolov3-tiny.cfg 是cfg中的yolov3的预训练参数
epochs 100 迭代训练数据集的次数


训练完成后会在weights文件夹下产生训练的模型，包括最好的训练模型和最后一次训练得到的模型还有每个每10次迭代产生的模型。这里我只保留了最好的以及最后一次的训练模型。

10. 测试模型

在完成了训练后，即可对训练的模型进行测试了，看一下这个其对佩戴口罩和未佩戴口罩者的预测效果
将要预测的实际照片放在data/sample文件夹下
在根目录下打开命令行，在其中输入如下命令：

python detect.py --data-cfg data/rbc.data --cfg cfg/yolov3-tiny.cfg --weights weights/best.pt

预测的结果如下：



如果需要对视频进行检测的话，则可以输入如下命令：

python detect.py --source file.mp4

11. 可能出现的bug

在操作的过程中难免会出现一些小毛病，这里呢，笔者已经为这个项目总结了几个容易出现的问题，请参考笔者的另一篇博客：项目可能出现的问题及其解决方法

项目源码和数据集

这里笔者将本项目的源代码和数据集都已经上传到github上，欢迎大家扩展，为保护我们的工作者减少工作压力。
链接： 源码和数据集

总结：

由于笔者知识能力有限，在描述上可能存在不准确的地方，还请谅解。
如遇到什么问题欢迎添加笔者qq:1017190168
进行讨论。
最后对一线的工作者们说一声你们辛苦了！！！！
祝愿疫情早日结束，祝祖国繁荣富强。