问题
由于新冠疫情的缘故，各地的公共场所都要求佩戴口罩，希望能够通过视频或者图片，识别出有没有佩戴口罩。
解决问题思路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 os3.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 target3.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.pth3.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=<AddBackward0>)
tensor(62.0787, grad_fn=<AddBackward0>)
tensor(54.0132, grad_fn=<AddBackward0>)
tensor(46.4962, grad_fn=<AddBackward0>)
tensor(42.1432, grad_fn=<AddBackward0>)
tensor(37.7047, grad_fn=<AddBackward0>)
tensor(32.0403, grad_fn=<AddBackward0>)
tensor(31.0030, grad_fn=<AddBackward0>)
tensor(30.6037, grad_fn=<AddBackward0>)
tensor(27.7224, grad_fn=<AddBackward0>)
tensor(28.1404, grad_fn=<AddBackward0>)
tensor(26.8357, grad_fn=<AddBackward0>)
tensor(27.4818, grad_fn=<AddBackward0>)
tensor(27.1225, grad_fn=<AddBackward0>)
tensor(23.8584, grad_fn=<AddBackward0>)
tensor(22.3662, grad_fn=<AddBackward0>)
tensor(24.2057, grad_fn=<AddBackward0>)
tensor(21.7948, grad_fn=<AddBackward0>)
tensor(21.4384, grad_fn=<AddBackward0>)
tensor(21.2375, grad_fn=<AddBackward0>)
tensor(20.6343, grad_fn=<AddBackward0>)
tensor(23.2005, grad_fn=<AddBackward0>)
tensor(18.7917, grad_fn=<AddBackward0>)
tensor(17.8378, grad_fn=<AddBackward0>)
tensor(19.9931, grad_fn=<AddBackward0>)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')