pytorch实现CNN卷积神经网络口罩识别

环境

• PyCharm Community Edition 2021.3.1
• Pytorch

代码实现

CNN卷积网络代码

import torch
from torch import nn
from torchvision import datasets
from torchvision import transforms
from torch.autograd import Variable
from torch.utils.data import DataLoader
from PIL import Image
import matplotlib.pyplot as plt


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(nn.Conv2d(3,16,3,padding=1),          # 第一个卷积层，输入通道数3，输出通道数16，卷积核大小3*3
                                    nn.ReLU(True),                        # 第一次卷积结果经过ReLU激活函数处理
                                    nn.MaxPool2d(kernel_size=2, stride=2) # 第一次池化，池化大小2*2，方式Max pooling
                                    )
        self.layer2 = nn.Sequential(nn.Conv2d(16,16,3,padding=1),
                                    nn.ReLU(True),
                                    nn.MaxPool2d(kernel_size=2, stride=2)
                                  )
        self.fc = nn.Sequential(nn.Linear(56 * 56 * 16, 128),# 第一个全连接层，线性连接，输入节点数56*56*16，输出节点数128
                                nn.ReLU(True),
                                nn.Linear(128, 64),# 第二个全连接层，线性连接，输入节点数128，输出节点数64
                                nn.ReLU(True),
                                nn.Linear(64, 2)# 第三个全连接层，线性连接，输入节点数64，输出节点数2
                                )
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.view(x.size()[0], -1)
        x = self.fc(x)
        return x

训练代码

from torch.utils.data import DataLoader
from torchvision import datasets
from CNN import CNN
import torchvision
from CNN import CNN
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import torchvision
from torchvision import transforms
from PIL import Image
from torch.autograd import Variable
device = torch.device('cuda')
# 0不戴口罩 1 戴口罩

trans = torchvision.transforms.Compose([torchvision.transforms.Resize((224, 224)),

                                        torchvision.transforms.ToTensor()])
root ="D:\desk\text"
data = datasets.ImageFolder(root="D:\desk\test", transform=trans)
loader = torch.utils.data.DataLoader(data, 1, True)

# 加载数据集
train_data = datasets.ImageFolder(root=r'D:\desk\test', transform=trans)
train_loader = torch.utils.data.DataLoader(train_data, 128, True)



learning_rate=0.01
criterion = nn.CrossEntropyLoss()
####################################
model = CNN()
if torch.cuda.is_available():
    print("CUDA is enable!")
    model = model.to(device)
####################################

epoches_num = 50

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
model.train(True)
for epoch in range(epoches_num):
    print('*' * 40)
    train_loss = 0.0
    train_acc = 0.0
    for step, data in enumerate(train_loader):
        inputs, label = data


        if torch.cuda.is_available():
            inputs = Variable(inputs).to(device)
            label = Variable(label).to(device)

        else:
            inputs = Variable(inputs)
            label = Variable

        optimizer.zero_grad()
        outputs = model(inputs)  # 网络前向传播

        criterion = criterion.to(device)
        loss = criterion(outputs, label)

        loss.backward()
        optimizer.step()


        train_loss += loss.item() * label.size(0)
        _, pred = outputs.max(1)

        num_correct = pred.eq(label).sum()
        accuracy = pred.eq(label).float().mean()
        train_acc += num_correct.item()
    torch.save(model, 'ooo.pt')
    print('Finish  {}  Loss: {:.6f}, Acc: {:.6f}'.format(epoch + 1, train_loss / len(train_data),
                                                         train_acc / len(train_data)))

视频实时监测

import cv2 as cv
from torch.utils.data import DataLoader
from torchvision import datasets
import torchvision
from CNN import CNN
import torch
import torch.nn as nn
import torch.optim
import matplotlib.pyplot as plt
import numpy
from torchvision import transforms
from PIL import Image
from torch.autograd import Variable

device = torch.device('cuda')
model = torch.load('maxcnn33.pt').to(device)
trans = torchvision.transforms.Compose([torchvision.transforms.Resize((224, 224)),
                                            torchvision.transforms.ToTensor()])
model.train(False)


def face_detect_demo(filename):
    face_cascade = cv.CascadeClassifier("D:\Anaconda\envs\pytorch-env\Lib\site-packages\cv2\data\haarcascade_frontalface_alt.xml")# 声明face_cascade为CascadeClassifier对象，它负责人脸检测。
    img = cv.imread(filename)
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray)


    for (x,y,w,h) in faces:
        cv.rectangle(img, (x, y), (x + w, y + h), color=(111, 222, 111), thickness=2)#框上
        #cv.imshow('result', img)
        cropped_image = img[y:y + h, x:x + w]  ##=裁剪后

        if len(cropped_image):

            cv.imwrite("D:\desk\text\d\face.jpg", cropped_image)
            image_path = Image.open("D:\desk\text\d\face.jpg")
            trans1 = trans(image_path)
            trans1 = torch.reshape(trans1, (1, 3, 224, 224)).to(device)
            OutPuts = model(trans1)
            print(OutPuts)
            _, pred = OutPuts.max(1)
            #print(pred)

            if pred == 1:
                print("with_mask")
                cv.putText(img, "with_mask", (y, x), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255))
            if pred == 0:
                print("without_mask")
                cv.putText(img, "without_mask", (y, x), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255))
            #cv.waitKey(0)
            cv.imshow('result', img)

        #elif(len(cropped_image)==0):

            #print("未检测到人脸")

        #cv.imshow('result', img)  # 显示
        cv.waitKey(10)#######!!!!!!!!!!!!!!!!!!!!!!



capture = cv.VideoCapture(0)
#cv.namedWindow("result", cv.WINDOW_AUTOSIZE)
while (True):
    ret, frame = capture.read()
    cv.imshow('1', frame)
    cv.imwrite("D:\desk\text\c\face.jpg", frame)
    path = "D:\desk\text\c\face.jpg"
    face_detect_demo(path)

cv.waitKey(0)
cv.destroyAllWindows()

由于训练样例为

此种全头像图片，所以在测试的时候，使用opencv的级联分类器获得人脸的坐标再将人脸进行裁剪，去除只保留头的部分，可以提高识别率 

cv.rectangle(img, (x, y), (x + w, y + h), color=(111, 222, 111), thickness=2)#框上
    
        cropped_image = img[y:y + h, x:x + w]  ##=裁剪后