基于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
进行讨论。
最后对一线的工作者们说一声你们辛苦了！！！！
祝愿疫情早日结束，祝祖国繁荣富强。
