基于深度学习的口罩识别与检测PyTorch实现
基于深度學(xué)習(xí)的口罩識別與檢測PyTorch實現(xiàn)
- 1. 設(shè)計思路
- 1.1 兩階段檢測器:先檢測人臉,然后將人臉進(jìn)行分類,戴口罩與不戴口罩。
- 1.2 一階段檢測器:直接訓(xùn)練口罩檢測器,訓(xùn)練樣本為人臉的標(biāo)注信息,戴口罩的,不帶口罩的。
- 2. 算法具體實現(xiàn)
- 2.1 兩階段口罩識別與檢測器設(shè)計思路概述
- 2.2 實現(xiàn)代碼分析
- 2.2.1 訓(xùn)練分類器并且保存模型
- 2.2.2 人臉檢測與剪裁,然后進(jìn)行分類。
- 3. 總結(jié)
- 4. 參考鏈接
1. 設(shè)計思路
- 已經(jīng)完成,兩階段口罩檢測的算法設(shè)計與實現(xiàn)。主要講解整體pipeline的設(shè)計,即:人臉檢測算法+口罩二分類CNN模型。
全文的代碼和技術(shù)講解,在這里可以下載。mask_detection代碼講解. - 一階段口罩檢測算法設(shè)計與實現(xiàn)。這一部分安排下一個博客講解。
1.1 兩階段檢測器:先檢測人臉,然后將人臉進(jìn)行分類,戴口罩與不戴口罩。
對于這種兩階段方法,是能快速實現(xiàn)的。首先用現(xiàn)有的人臉檢測算法,直接對圖像進(jìn)行人臉檢測,然后將檢測的每一個人臉,單獨切割出來,進(jìn)行是否戴口罩的二分類。
所以只需要訓(xùn)練一個CNN網(wǎng)絡(luò),包含兩類圖像,mask和without_maks。前端接一個face detector,就能應(yīng)對大多數(shù)情況了。單獨的講解文檔鏈接:口罩檢測思路講解文檔。
1.2 一階段檢測器:直接訓(xùn)練口罩檢測器,訓(xùn)練樣本為人臉的標(biāo)注信息,戴口罩的,不帶口罩的。
先對圖像中的人臉進(jìn)行標(biāo)注,包括了戴口罩的,和不帶口罩的兩個label的目標(biāo)。然后進(jìn)行標(biāo)注。直接對人臉檢測模型進(jìn)行transfer learning,微調(diào)一下模型,就能檢測了。
直接對戴口罩和不帶口罩的人臉進(jìn)行標(biāo)注就可以了。
類似這種標(biāo)注方法:
2. 算法具體實現(xiàn)
2.1 兩階段口罩識別與檢測器設(shè)計思路概述
-
用數(shù)據(jù)集訓(xùn)練一個CNN二分類器。
其中的數(shù)據(jù)為兩類,mask 和without_mask。可以用現(xiàn)成的模型參數(shù),在ImageNet上預(yù)訓(xùn)練好的weights,進(jìn)行finetune。本方法采用的是ResNet18 為backbone,來進(jìn)行訓(xùn)練。詳細(xì)代碼如下。其中,mask的數(shù)據(jù)集來這里: Mask Dataset
如果這個數(shù)據(jù)集鏈接失效,用這一個:Mask Dataset 備份地址
注意,下載數(shù)據(jù)集是google drive,科學(xué)上網(wǎng)。
對訓(xùn)練好的模型,進(jìn)行權(quán)重參數(shù)的保存,用于第二階段的分類任務(wù)。保存為:將訓(xùn)練好的模型保存在pth中,然后在evaluation中,或者在Inference中,直接使用這個模型參數(shù)。## 將訓(xùn)練好的模型保存在pth中,然后在evaluation中,或者在Inference中,直接使用這個模型參數(shù)。 torch.save(model_ft, 'finetuned_model_resnet18.pth') -
使用人臉檢測器,檢測人臉,切割人臉。這里如果有必要,需要重點關(guān)注戴口罩的人臉的檢測效果。
-
對切割的人臉,用上述訓(xùn)練好的分類器,對人臉進(jìn)行分類:戴口罩和不帶口罩的。
-
下面的訓(xùn)練分類器的代碼,和人臉檢測與分類代碼,都在這個鏈接里面下載。mask_detection代碼詳細(xì)講解。
2.2 實現(xiàn)代碼分析
2.2.1 訓(xùn)練分類器并且保存模型
from __future__ import print_function, divisionimport torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
import PIL.ImageOps
import requests
from PIL import Image# Load Data
# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {'train': transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),'val': transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
}data_dir = 'data/mask'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),data_transforms[x])for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,shuffle=True, num_workers=4)for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classesdevice = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#Visualize a few imagesdef imshow(inp, title=None):"""Imshow for Tensor."""inp = inp.numpy().transpose((1, 2, 0))mean = np.array([0.485, 0.456, 0.406])std = np.array([0.229, 0.224, 0.225])inp = std * inp + meaninp = np.clip(inp, 0, 1)plt.imshow(inp)if title is not None:plt.title(title)plt.pause(0.001) # pause a bit so that plots are updated# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))# Make a grid from batch
out = torchvision.utils.make_grid(inputs)imshow(out, title=[class_names[x] for x in classes])def train_model(model, criterion, optimizer, scheduler, num_epochs=25):since = time.time()best_model_wts = copy.deepcopy(model.state_dict())best_acc = 0.0for epoch in range(num_epochs):print('Epoch {}/{}'.format(epoch, num_epochs - 1))print('-' * 10)# Each epoch has a training and validation phasefor phase in ['train', 'val']:if phase == 'train':model.train() # Set model to training modeelse:model.eval() # Set model to evaluate moderunning_loss = 0.0running_corrects = 0# Iterate over data.for inputs, labels in dataloaders[phase]:inputs = inputs.to(device)labels = labels.to(device)# zero the parameter gradientsoptimizer.zero_grad()# forward# track history if only in trainwith torch.set_grad_enabled(phase == 'train'):outputs = model(inputs)_, preds = torch.max(outputs, 1)loss = criterion(outputs, labels)# backward + optimize only if in training phaseif phase == 'train':loss.backward()optimizer.step()# statisticsrunning_loss += loss.item() * inputs.size(0)running_corrects += torch.sum(preds == labels.data)if phase == 'train':scheduler.step()epoch_loss = running_loss / dataset_sizes[phase]epoch_acc = running_corrects.double() / dataset_sizes[phase]print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))# deep copy the modelif phase == 'val' and epoch_acc > best_acc:best_acc = epoch_accbest_model_wts = copy.deepcopy(model.state_dict())print()time_elapsed = time.time() - sinceprint('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))print('Best val Acc: {:4f}'.format(best_acc))# load best model weightsmodel.load_state_dict(best_model_wts)return model#Visualizing the model predictionsdef visualize_model(model, num_images=6):was_training = model.trainingmodel.eval()images_so_far = 0fig = plt.figure()with torch.no_grad():for i, (inputs, labels) in enumerate(dataloaders['val']):inputs = inputs.to(device)labels = labels.to(device)outputs = model(inputs)_, preds = torch.max(outputs, 1)for j in range(inputs.size()[0]):images_so_far += 1ax = plt.subplot(num_images//2, 2, images_so_far)ax.axis('off')ax.set_title('predicted: {}'.format(class_names[preds[j]]))imshow(inputs.cpu().data[j])if images_so_far == num_images:model.train(mode=was_training)returnmodel.train(mode=was_training)if __name__ == "__main__":#Finetuning the convnet
#Load a pretrained model and reset final fully connected layer.model_ft = models.resnet18(pretrained=True)num_ftrs = model_ft.fc.in_features# Here the size of each output sample is set to 2.# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).model_ft.fc = nn.Linear(num_ftrs, 2)model_ft = model_ft.to(device)criterion = nn.CrossEntropyLoss()# Observe that all parameters are being optimizedoptimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)# Decay LR by a factor of 0.1 every 7 epochsexp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,num_epochs=25)torch.save(model_ft, 'finetuned_model_resnet18.pth')visualize_model(model_ft)
這是一些數(shù)據(jù)集的可視化圖:
訓(xùn)練后的分類結(jié)果展示:
2.2.2 人臉檢測與剪裁,然后進(jìn)行分類。
這里我們采用MTCNN進(jìn)行人臉檢測。
import cv2
from mtcnn.core.detect import create_mtcnn_net, MtcnnDetector
from mtcnn.core.vision import vis_faceif __name__ == '__main__':### part1: face detectionpnet, rnet, onet = create_mtcnn_net(p_model_path="./original_model/pnet_epoch.pt", r_model_path="./original_model/rnet_epoch.pt", o_model_path="./original_model/onet_epoch.pt", use_cuda=False)mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, onet=onet, min_face_size=24)img = cv2.imread("./s_l.jpg")img_bg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)bboxs, landmarks = mtcnn_detector.detect_face(img)save_name = 'r_1.jpg'## visulaze the face detection effect. vis_face(img_bg,bboxs,None, save_name)### part2: mask classification.## load classification model.mask_model = torch.load('finetuned_model_resnet18.pth',map_location='cpu')mask_list = []for i in range(bboxs.shape[0]):bbox = bboxs[i, :4]print(img.shape)print(bbox)## 這里主要是人臉的剪裁,然后數(shù)據(jù)的轉(zhuǎn)換,mat 格式轉(zhuǎn)換為 numpy。cropImg1 = img[int(bbox[1]):int(bbox[3]),int(bbox[0]):int(bbox[2])]cv2.imshow('face',cropImg1)cv2.waitKey(0)## 剪裁人臉的格式轉(zhuǎn)換PIL_image = Image.fromarray(cropImg1)face_img = data_transforms['val'](PIL_image) face_img = face_img.to(device).unsqueeze(0)output = mask_model (face_img)_, pred = torch.max(output, 1)mask_list.append(class_names[pred.item()])print(class_names[pred.item()])### 可視化的一些操作,包括人臉的rectangle以及putText到圖像上。cv2.rectangle(img, (int(bbox[0]),int(bbox[1])), (int(bbox[2]),int(bbox[3])), (0,255,0), 4)font = cv2.FONT_HERSHEY_SIMPLEXcv2.putText(img, class_names[pred.item()], (int(bbox[0]),int(bbox[1])), font, 1, (0,0,255), 2)cv2.imshow('original',img)# cv2.imshow('face',cropImg1)cv2.waitKey(0)cv2.imwrite('result_new.jpg', img)img_result = Image.open('result_new.jpg')plt.imshow(img_result)
算法的效果圖
- 先是人臉檢測效果:
- 識別分類后的效果:
- 大規(guī)模實際效果圖
用比較強(qiáng)的人臉檢測算法,能夠檢測出來人臉,然后在進(jìn)行分類。
3. 總結(jié)
由于采用兩階段的方法,依賴于人臉檢測算法的魯棒性,如果人臉檢測算法,對戴口罩之后的遮擋,大面積遮擋,導(dǎo)致人臉特征減少,導(dǎo)致不能檢出人臉。這時候,第二階段的分類,是沒有意義的。因為第一步就檢測不出人臉來,性能差主要是第一階段的能力導(dǎo)致的,而對于分類這一塊,分類器沒什么大問題。因此主要的精力放在人臉檢測器的優(yōu)化上來。由于原始數(shù)據(jù)集沒有戴口罩的人臉,因此需要對數(shù)據(jù)分布進(jìn)行調(diào)整。關(guān)于這些缺點和對應(yīng)解決方案的描述文檔,在這里。口罩檢測描述與解決方案的文檔。
下面的圖像就是戴口罩之后檢測不出例子,當(dāng)口罩太大,遮住2/3的人臉之后,人臉檢測算法檢測不出來。
一些原因描述:
解決方法:
- 重新訓(xùn)練人臉檢測算法,將戴口罩,不帶口罩的人臉,全部標(biāo)注出來,都標(biāo)注為人臉,然后訓(xùn)練模型,加強(qiáng)其對戴口罩,遮擋住嘴巴鼻子的這一類人臉的檢出率。重點關(guān)注這種戴口罩人臉的檢測效果。 主要訓(xùn)練人臉檢測模型,在普通人臉的數(shù)據(jù)基礎(chǔ)上,增強(qiáng)人臉檢測性能,提升戴口罩人臉檢測的魯棒性。
- 當(dāng)然是直接使用人臉檢測算法,將人臉標(biāo)注為兩類目標(biāo),戴口罩的,不帶口罩的。直接在人臉檢測的基礎(chǔ)上,進(jìn)行兩類目標(biāo)的直接檢測。
其中絕大部分的問題和解決方法,都在我給的技術(shù)文檔,ppt,或者下面的鏈接中,都描述了的,所以認(rèn)真看一下我給的文件。mask_detection文檔與代碼講解。
- 已經(jīng)完成,兩階段口罩檢測的算法設(shè)計與實現(xiàn)。
- TODO:一階段口罩檢測算法設(shè)計與實現(xiàn)。
4. 參考鏈接
- https://www.analyticsvidhya.com/blog/2020/08/how-to-build-a-face-mask-detector-using-retinanet-model/
- https://github.com/chandrikadeb7/Face-Mask-Detection
- https://towardsdatascience.com/covid-19-face-mask-detection-using-tensorflow-and-opencv-702dd833515b
- https://www.mygreatlearning.com/blog/real-time-face-detection/
- https://data-flair.training/blogs/face-mask-detection-with-python/
- https://www.pyimagesearch.com/2020/05/04/covid-19-face-mask-detector-with-opencv-keras-tensorflow-and-deep-learning/
- https://www.ideas2it.com/blogs/face-mask-detector-using-deep-learning-pytorch-and-computer-vision-opencv/
如果有用,記得點贊👍加收藏哦。!!!!
總結(jié)
以上是生活随笔為你收集整理的基于深度学习的口罩识别与检测PyTorch实现的全部內(nèi)容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 康复新液多少钱啊?
- 下一篇: 用Macbook-苹果系统写代码出现显示