FastRCNN实现识别口罩

问题

由于新冠疫情的缘故，各地的公共场所都要求佩戴口罩，希望能够通过视频或者图片，识别出有没有佩戴口罩。

解决问题思路FastRCNN

a. 在图像中确定N个候选框

b. 对于每个候选框内图像块，使用深度网络提取特征

c. 对候选框中提取出的特征，使用分类器判别是否属于一个特定类

d. 对于属于某一特征的候选框，用回归器进一步调整其位置

[1]图片地址

训练集地址 https://www.kaggle.com/andrewmvd/face-mask-detection

网盘地址（密码: w7ti）

[2]代码结构

代码结构如下

[3]构造训练模型 face_mask_train.py

3.1 引入依赖

from bs4 import BeautifulSoup

import torchvision

from torchvision import transforms, datasets, models

import torch

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

from PIL import Image

import matplotlib.pyplot as plt

import matplotlib.patches as patches

import os

3.2 数据结构化函数

def generate_box(obj):
    xmin = int(obj.find('xmin').text)
    ymin = int(obj.find('ymin').text)
    xmax = int(obj.find('xmax').text)
    ymax = int(obj.find('ymax').text)
    return [xmin, ymin, xmax, ymax]

def generate_label(obj):
    if obj.find('name').text == "with_mask":
        return 1
    elif obj.find('name').text == "mask_weared_incorrect":
        return 2
    return 0


def generate_target(image_id, file):
    with open(file) as f:
        data = f.read()
        soup = BeautifulSoup(data, 'xml')
        objects = soup.find_all('object')4
        boxes = []
        labels = []
        for i in objects:
            boxes.append(generate_box(i))
            labels.append(generate_label(i))

        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(labels, dtype=torch.int64)
        img_id = torch.tensor([image_id])

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = img_id
        return target

3.3 训练模型地址

imgs = list(sorted(os.listdir("./face-mask/images/")))

labels = list(sorted(os.listdir("./face-mask/annotations/")))

3.4 构造口罩模型类

class MaskDataset(object):

    def __init__(self, transforms):
        self.transforms = transforms
        self.imgs = list(sorted(os.listdir("./face-mask/images/")))

    def __getitem__(self, idx):
        file_image = 'maksssksksss'+ str(idx) + '.png'
        file_label = 'maksssksksss'+ str(idx) + '.xml'
        img_path = os.path.join("./face-mask/images/", file_image)
        label_path = os.path.join("./face-mask/annotations/", file_label)
        img = Image.open(img_path).convert("RGB")
        #Generate Label
        target = generate_target(idx, label_path)

        if self.transforms is not None:
            img = self.transforms(img)

        return img, target

    def __len__(self):
        return len(self.imgs)

3.5 数据集合构造

data_transform = transforms.Compose([

        #函数接受PIL Image或numpy.ndarray，将其先由HWC转置为CHW格式，再转为float后每个像素除以255.
        transforms.ToTensor(),
    ])


def collate_fn(batch):
    return tuple(zip(*batch))

dataset = MaskDataset(data_transform)
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=4, collate_fn=collate_fn)
torch.cuda.is_available()

3.6 构造模型

def get_model_instance_segmentation(num_classes):
    # load an instance segmentation model pre-trained pre-trained on COCO
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model
model = get_model_instance_segmentation(3)

下载训练模型

Downloading: "https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth" to /root/.cache/torch/checkpoints/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth

3.7 训练模型

#我的电脑不支持，大家可以自己看是否支持cuda
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
for imgs, annotations in data_loader:
    imgs = list(img.to(device) for img in imgs)
    annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
    print(annotations)
    break

num_epochs = 25
model.to(device)
# parameters
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
                                momentum=0.9, weight_decay=0.0005)

len_dataloader = len(data_loader)
for epoch in range(num_epochs):
    model.train()
    i = 0    
    epoch_loss = 0
    for imgs, annotations in data_loader:
        i += 1
        imgs = list(img.to(device) for img in imgs)
        annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
        loss_dict = model([imgs[0]], [annotations[0]])
        losses = sum(loss for loss in loss_dict.values())        

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        epoch_loss += losses
    print(epoch_loss)

时间会比较长，我的电脑性能不好，差不多用了一天的时间，一定要耐心等待

[{'boxes': tensor([[ 79., 105., 109., 142.],
        [185., 100., 226., 144.],
        [325.,  90., 360., 141.]]), 'labels': tensor([0, 1, 0]), 'image_id': tensor([0])}, {'boxes': tensor([[321.,  34., 354.,  69.],
        [224.,  38., 261.,  73.],
        [299.,  58., 315.,  81.],
        [143.,  74., 174., 115.],
        [ 74.,  69.,  95.,  99.],
        [191.,  67., 221.,  93.],
        [ 21.,  73.,  44.,  93.],
        [369.,  70., 398.,  99.],
        [ 83.,  56., 111.,  89.]]), 'labels': tensor([1, 1, 1, 1, 1, 1, 1, 1, 0]), 'image_id': tensor([1])}, {'boxes': tensor([[ 68.,  42., 105.,  69.],
        [154.,  47., 178.,  74.],
        [238.,  34., 262.,  69.],
        [333.,  31., 366.,  65.]]), 'labels': tensor([1, 1, 1, 2]), 'image_id': tensor([2])}, {'boxes': tensor([[ 52.,  53.,  73.,  76.],
        [ 72.,  53.,  92.,  75.],
        [112.,  51., 120.,  68.],
        [155.,  60., 177.,  83.],
        [189.,  59., 210.,  80.],
        [235.,  57., 257.,  78.],
        [289.,  60., 309.,  83.],
        [313.,  68., 333.,  90.],
        [351.,  35., 364.,  59.]]), 'labels': tensor([1, 1, 1, 1, 1, 1, 1, 1, 1]), 'image_id': tensor([3])}]

/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  ../c10/core/TensorImpl.h:1156.)
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
tensor(87.1121, grad_fn=)
tensor(62.0787, grad_fn=)
tensor(54.0132, grad_fn=)
tensor(46.4962, grad_fn=)
tensor(42.1432, grad_fn=)
tensor(37.7047, grad_fn=)
tensor(32.0403, grad_fn=)
tensor(31.0030, grad_fn=)
tensor(30.6037, grad_fn=)
tensor(27.7224, grad_fn=)
tensor(28.1404, grad_fn=)
tensor(26.8357, grad_fn=)
tensor(27.4818, grad_fn=)
tensor(27.1225, grad_fn=)
tensor(23.8584, grad_fn=)
tensor(22.3662, grad_fn=)
tensor(24.2057, grad_fn=)
tensor(21.7948, grad_fn=)
tensor(21.4384, grad_fn=)
tensor(21.2375, grad_fn=)
tensor(20.6343, grad_fn=)
tensor(23.2005, grad_fn=)
tensor(18.7917, grad_fn=)
tensor(17.8378, grad_fn=)
tensor(19.9931, grad_fn=)

3.8 模型结果

for imgs, annotations in data_loader:
        imgs = list(img.to(device) for img in imgs)
        annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
        break
model.eval()
preds = model(imgs)
preds

def plot_image(img_tensor, annotation):
    fig,ax = plt.subplots(1)
    img = img_tensor.cpu().data

    # Display the image
    ax.imshow(img.permute(1, 2, 0))

    for box in annotation["boxes"]:
        xmin, ymin, xmax, ymax = box

        # Create a Rectangle patch
        rect = patches.Rectangle((xmin,ymin),(xmax-xmin),(ymax-ymin),linewidth=1,edgecolor='r',facecolor='none')

        # Add the patch to the Axes
        ax.add_patch(rect)
    plt.show()

print("Prediction")
plot_image(imgs[2], preds[2])
print("Target")
plot_image(imgs[2], annotations[2])

验证结果

Prediction

Target

3.9 模型保存！！训练这个太久了，赶紧保存下

torch.save(model.state_dict(),'model.pt')