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

2022-05-16 18:15:55 浏览数 (1)

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

1. 开发环境

笔者的开发环境如下:

ubuntu16.04

PyTorch 1.1.0

anaconda

opencv-python

tqdm

matplotlib

pycocotools

2. 收集数据集

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

代码语言:txt复制
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是对获取的数据进行命名整理的。代码如下:

代码语言:python代码运行次数:0复制
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的代码如下:

代码语言:python代码运行次数:0复制
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的代码如下:

代码语言:python代码运行次数:0复制
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.jpgn' % (image_id))
        convert_annotation(image_id)
    list_file.close()

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

在这里插入图片描述在这里插入图片描述

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

在这里插入图片描述在这里插入图片描述

接着还要配置两个文件

在data文件下新建rbc.data,配置内容如下:

代码语言:python代码运行次数:0复制
classes=2
train=data/train.txt
valid=data/test.txt
names=data/rbc.names
backup=backup/
eval=coco

值得注意的是这里是对人是否佩戴口罩进行识别的,即分为两类佩戴口罩和为佩戴口罩。所以这里需要将classes设置为2.

再在data文件下新建rbc.names,配置内容如下:

代码语言:shell复制
Masking
unMasked

8. 修改配置文件

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

代码语言:shell复制
[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. 训练数据集

完成了上述操作后,即可进行数据集的训练操作了。

在根目录下打开命令行,输入如下命令:

代码语言:shell复制
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文件夹下

在根目录下打开命令行,在其中输入如下命令:

代码语言:shell复制
python detect.py --data-cfg data/rbc.data --cfg cfg/yolov3-tiny.cfg --weights weights/best.pt

预测的结果如下:

在这里插入图片描述在这里插入图片描述
在这里插入图片描述在这里插入图片描述
在这里插入图片描述在这里插入图片描述在这里插入图片描述在这里插入图片描述

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

代码语言:shell复制
python detect.py --source file.mp4

11. 可能出现的bug

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

项目源码和数据集

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

链接: 源码和数据集

总结:

由于笔者知识能力有限,在描述上可能存在不准确的地方,还请谅解。

最后对一线的工作者们说一声你们辛苦了!!!!

祝愿疫情早日结束,祝祖国繁荣富强。

image.pngimage.png

欢迎关注公众号:陶陶name

0 人点赞