其他
AI教程|用滴滴云Notebook笔记本服务实现Yolo3口罩检测
在本教程中，您将学习如何使用经典的one-stage目标检测网络Yolo v3来实现口罩检测，关于Yolo v3的资料可以阅读paper。本教程所有工程文件可在滴滴云S3存储服务下载。
滴滴云Notebook笔记本服务集成了CUDA、CuDNN、Python、TensorFlow、Pytorch、MxNet、Keras等深度学习框架，无需用户自己安装。
一、购买Notebook服务
注册滴滴云并实名认证后可购买Notebook服务。
进入控制台Notebook页面，单击创建Notebook实例按钮。
选择基础配置：
选择 付费方式:当前仅支持按时长。
选择 可用区:选择靠近您客户的地域，包括广州1、2区。
选择 配置规格:根据需要的CPU、GPU、显卡和内存，选择相关配置。
选择 镜像:提供了Jupyter Notebook镜像和Jupyter Lab镜像，这里选择>jupyter-lab-v1。
设置 系统盘:根据需求选择系统盘的大小，设置范围为80GB - 500GB。
设置名称和标签
输入 Notebook名称。
输入 标签键以及键值，单击添加按钮，可添加多个标签。
访问Notebook
进入我的Notebook页面，在操作列单击打开Notebook。
进入Notebook详情页面，单击打开Notebook。
二、调试代码
import cv2
import math
import matplotlib.pyplot as plt
import numpy as np
import os
import random
import time
import torch
import torchvision
import torch.nn as nn
import torch.nn.init as init
import torch.optim as optim
import xml.etree.ElementTree as ET
from torch.utils.data import Dataset, DataLoader下载口罩检测数据集并上传到Notebook服务器，这里我们以AIZOO开源数据集为例，下载地址：
https://pan.baidu.com/s/1nsQf_Py5YyKm87-8HiyJeQ ，提取码：eyfz
大文件上传可能需要用到滴滴云S3存储服务。为了便于用户下载，我们提前上传到公共数据集对象存储S3，执行以下shell命令即可，如自行上传可以跳过。
!wget https://dataset-public.s3.didiyunapi.com/detection/人脸口罩检测/part1.tgz
!wget https://dataset-public.s3.didiyunapi.com/detection/人脸口罩检测/part2.tgz
!wget https://dataset-public.s3.didiyunapi.com/detection/人脸口罩检测/val.tgz
!tar -zxf part1.tgz
!tar -zxf part2.tgz
!tar -zxf val.tgz
--2020-03-05 14:59:25-- https://dataset-public.s3.didiyunapi.com/detection/%E4%BA%BA%E8%84%B8%E5%8F%A3%E7%BD%A9%E6%A3%80%E6%B5%8B/part1.tgz
Resolving dataset-public.s3.didiyunapi.com (dataset-public.s3.didiyunapi.com)... 125.94.54.9
Connecting to dataset-public.s3.didiyunapi.com (dataset-public.s3.didiyunapi.com)|125.94.54.9|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 270682151 (258M) [application/gzip]
Saving to: ‘part1.tgz’
100%[======================================>] 270,682,151 6.15MB/s in 42s
2020-03-05 15:00:06 (6.18 MB/s) - ‘part1.tgz’ saved [270682151/270682151]
--2020-03-05 15:00:07-- https://dataset-public.s3.didiyunapi.com/detection/%E4%BA%BA%E8%84%B8%E5%8F%A3%E7%BD%A9%E6%A3%80%E6%B5%8B/part2.tgz
Resolving dataset-public.s3.didiyunapi.com (dataset-public.s3.didiyunapi.com)... 125.94.54.9
Connecting to dataset-public.s3.didiyunapi.com (dataset-public.s3.didiyunapi.com)|125.94.54.9|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 337432016 (322M) [application/gzip]
Saving to: ‘part2.tgz’
100%[======================================>] 337,432,016 6.41MB/s in 53s
2020-03-05 15:00:59 (6.10 MB/s) - ‘part2.tgz’ saved [337432016/337432016]
--2020-03-05 15:01:00-- https://dataset-public.s3.didiyunapi.com/detection/%E4%BA%BA%E8%84%B8%E5%8F%A3%E7%BD%A9%E6%A3%80%E6%B5%8B/val.tgz
Resolving dataset-public.s3.didiyunapi.com (dataset-public.s3.didiyunapi.com)... 125.94.54.9
Connecting to dataset-public.s3.didiyunapi.com (dataset-public.s3.didiyunapi.com)|125.94.54.9|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 184383116 (176M) [application/gzip]
Saving to: ‘val.tgz’
100%[======================================>] 184,383,116 6.25MB/s in 29s
2020-03-05 15:01:28 (6.12 MB/s) - ‘val.tgz’ saved [184383116/184383116]加载数据，这里我们需要编写自定义数据集的Dataset类型
VOC_CLASSES = ('face', 'face_mask')
class AnnotationTransform(object):
def __init__(self, class_to_ind=None, keep_difficult=True):
self.class_to_ind = class_to_ind or dict(
zip(VOC_CLASSES, range(len(VOC_CLASSES))))
self.keep_difficult = keep_difficult
def __call__(self, target):
res = np.empty((0,5))
for obj in target.iter('object'):
difficult = int(obj.find('difficult').text) == 1
if not self.keep_difficult and difficult:
continue
name = obj.find('name').text.lower().strip()
bbox = obj.find('bndbox')
pts = ['xmin', 'ymin', 'xmax', 'ymax']
bndbox = []
for i, pt in enumerate(pts):
cur_pt = int(bbox.find(pt).text) - 1
bndbox.append(cur_pt)
label_idx = self.class_to_ind[name]
bndbox.append(label_idx)
res = np.vstack((res, bndbox)) # [xmin, ymin, xmax, ymax, label_ind]
return res # [[xmin, ymin, xmax, ymax, label_ind], ... ]
def preproc_for_test(image, input_size, mean, std):
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
interp_method = interp_methods[random.randrange(5)]
image = cv2.resize(image, input_size, interpolation=interp_method)
image = image.astype(np.float32)
image = image[:, :, ::-1]
image /= 255.
if mean is not None:
image -= mean
if std is not None:
image /= std
return image.transpose(2, 0, 1)
class TrainTransform(object):
def __init__(self, rgb_means=None, std=None, max_labels=50):
self.means = rgb_means
self.std = std
self.max_labels = max_labels
def __call__(self, image, targets, img_size):
boxes = targets[:, :4].copy() # Nx4
labels = targets[:, 4].copy()
if len(boxes) == 0:
targets = np.zeros((self.max_labels, 5), dtype=np.float32)
image = preproc_for_test(image, img_size, self.means, self.std)
image = np.ascontiguousarray(image, dtype=np.float32)
return torch.from_numpy(image), torch.from_numpy(targets)
height, width, _ = image.shape
boxes_o = targets[:, :4]
labels = targets[:, 4]
b_x_o = (boxes_o[:, 2] + boxes_o[:, 0]) * .5
b_y_o = (boxes_o[:, 3] + boxes_o[:, 1]) * .5
b_w_o = (boxes_o[:, 2] - boxes_o[:, 0]) * 1.
b_h_o = (boxes_o[:, 3] - boxes_o[:, 1]) * 1.
boxes_o[:, 0] = b_x_o
boxes_o[:, 1] = b_y_o
boxes_o[:, 2] = b_w_o
boxes_o[:, 3] = b_h_o
# resize
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
interp_method = interp_methods[random.randrange(5)]
image_t = cv2.resize(image, img_size, interpolation=interp_method)
boxes = boxes_o
boxes[:, 0::2] /= width
boxes[:, 1::2] /= height
boxes[:, 0::2] *= img_size[0]
boxes[:, 1::2] *= img_size[1]
image_t = preproc_for_test(image_t, img_size, self.means, self.std)
labels = np.expand_dims(labels, 1)
targets_t = np.hstack((labels, boxes))
padded_labels = np.zeros((self.max_labels, 5))
padded_labels[range(len(targets_t))[:self.max_labels]] = targets_t[:self.max_labels]
padded_labels = np.ascontiguousarray(padded_labels, dtype=np.float32)
image_t = np.ascontiguousarray(image_t, dtype=np.float32)
return torch.from_numpy(image_t), torch.from_numpy(padded_labels)
# 数据集类型定义
class VOCDetection(Dataset):
def __init__(self, root, preproc=None, target_transform=AnnotationTransform(), img_size=(416, 416), split='train'):
super().__init__()
self.root = root
self.preproc = preproc
self.target_transform = target_transform
self.img_size = img_size
self._annopath = os.path.join('%s', 'Annotations', '%s.xml')
self._imgpath = os.path.join('%s', 'JPEGImages', '%s.jpg')
self._classes = VOC_CLASSES
self._year = '2012' # options: '2007', which is related to eval protocol
self.item_container = set()
if split == 'train':
for folder in ['part1', 'part2']:
for item in os.listdir(os.path.join(self.root, folder)):
self.item_container.add(os.path.join(self.root, folder, item[:-4]))
else:
for folder in ['val']:
for item in os.listdir(os.path.join(self.root, folder)):
self.item_container.add(os.path.join(self.root, folder, item[:-4]))
self.item_container = list(self.item_container)
def __getitem__(self, index):
item = self.item_container[index]
target = ET.parse(item+'.xml').getroot()
img = cv2.imread(item+'.jpg')
# img = Image.open(self._imgpath % img_id).convert('RGB')
height, width, _ = img.shape
if self.target_transform is not None:
target = self.target_transform(target)
if self.preproc is not None:
img, target = self.preproc(img, target, self.img_size)
img_info = (width, height)
return img, target, img_info, item
def __len__(self):
return len(self.item_container)
dataset = VOCDetection(root='./', preproc=TrainTransform(),split='train')定义模型，对模型参数进行初始化操作，这里需要导入自定义库yolo，可在滴滴云S3存储服务下载。
from yolo import YOLOv3
model = YOLOv3(num_classes = len(VOC_CLASSES))
def init_yolo(M):
for m in M.modules():
if isinstance(m, nn.Conv2d):
init.kaiming_normal_(m.weight, a=0.1, mode='fan_in')
if m.bias is not None:
init.zeros_(m.bias)
elif isinstance(m, nn.BatchNorm2d):
init.ones_(m.weight)
init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
init.normal_(m.weight, 0, 0.01)
init.zeros_(m.bias)
m.state_dict()[key][...] = 0
model.apply(init_yolo)
model.train()
torch.backends.cudnn.benchmark = True
device = torch.device("cuda")
model = model.to(device)现在我们使用Adam算法来训练模型
# 在训练之前我们先定义一些超参数
batch_size = 8 # 每一批次训练大小，不宜太小，大小受GPU显存限制
base_lr = 0.0001 # 基准学习率
warmup_epochs = 10 # 学习率逐渐增加到base_lr的epoch
epochs = 70 # 总共训练epoch数
save_interval = 10 # 保存模型的epoch间隔
steps = [50, 60] # 学习率减少的epoch
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
optimizer = optim.Adam(model.parameters(), lr=base_lr, weight_decay=0.0005)
epoch_size = len(dataset) // (batch_size*1)
epoch = 1
def set_lr(tmp_lr):
for param_group in optimizer.param_groups:
param_group['lr'] = tmp_lr
while epoch < epochs+1:
print('\n[Epoch {} started]'.format(epoch))
for iter_i, (imgs, targets, _, _) in enumerate(dataloader):
start = time.time()
if epoch % save_interval == 0:
torch.save(model.state_dict(), 'yolov3_mask_detection_{}.pth'.format(epoch))
# 更新学习率
if epoch < warmup_epochs:
tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (warmup_epochs*epoch_size), 1)
set_lr(tmp_lr)
elif epoch == warmup_epochs:
tmp_lr = base_lr
set_lr(tmp_lr)
elif epoch in steps and iter_i == 0:
tmp_lr = tmp_lr * 0.1
set_lr(tmp_lr)
optimizer.zero_grad()
imgs = imgs.to(device).to(torch.float32)
targets = targets.to(device).to(torch.float32)
loss_dict = model(imgs, targets, epoch)
loss = sum(loss for loss in loss_dict['losses'])
loss.backward()
optimizer.step()
end = time.time()
if iter_i % 1 == 0:
# 打印训练过程信息
print('\r[Epoch %d/%d][Iter %d/%d][LR %.6f]'
'[Loss: l1 %.2f, conf %.6f, cls %.6f][Time: %.2f s]......'
% (epoch, epochs, iter_i+1, epoch_size, tmp_lr,
sum(l1_loss for l1_loss in loss_dict['l1_losses']).item(),
sum(conf_loss for conf_loss in loss_dict['conf_losses']).item(),
sum(cls_loss for cls_loss in loss_dict['cls_losses']).item(), end-start),
end='')
epoch += 1
torch.save(model.state_dict(), 'yolov3_mask_detection_final.pth'.format(epoch))[Epoch 1 started]
[Epoch 1/70][Iter 765/765][LR 0.000020][Loss: l1 4.36, conf 1747.822632, cls 1.573278][Time: 0.61 s]........
[Epoch 2 started]
[Epoch 2/70][Iter 765/765][LR 0.000030][Loss: l1 5.01, conf 781.659180, cls 1.723707][Time: 0.63 s].........
[Epoch 3 started]
[Epoch 3/70][Iter 765/765][LR 0.000040][Loss: l1 18.94, conf 331.378754, cls 5.523039][Time: 0.69 s].......
[Epoch 4 started]
[Epoch 4/70][Iter 78/765][LR 0.000041][Loss: l1 6.72, conf 293.010193, cls 1.820950][Time: 0.66 s].......
KeyboardInterrupt:得到训练的模型之后可以开始测试
class ValTransform(object):
def __init__(self, rgb_means=None, std=None, swap=(2, 0, 1)):
self.means = rgb_means
self.swap = swap
self.std = std
def __call__(self, img, res, input_size):
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
interp_method = interp_methods[0]
img = cv2.resize(np.array(img), input_size, interpolation=interp_method).astype(np.float32)
img = img[:, :, ::-1]
img /= 255.
if self.means is not None:
img -= self.means
if self.std is not None:
img /= self.std
img = img.transpose(self.swap)
img = np.ascontiguousarray(img, dtype=np.float32)
return torch.from_numpy(img), torch.zeros(1, 5)
transform = ValTransform()
im = cv2.imread("val/test_00000760.jpg") # 输入的图片
ori_im = im.copy()
height, width, _ = im.shape
test_size = (416, 416)
im_input, _ = transform(im, None, test_size)
im_input = im_input.to(device).type(torch.float32).unsqueeze(0)
model.load_state_dict(torch.load('yolov3_mask_detection_final.pth')) # 加载训练权重
device = torch.device("cuda")
model = model.to(device)
model.eval()
outputs = model(im_input)对模型输出做后处理，除去置信度较低的bbox，并利用非极大抑制（NMS）去除同类型IoU较大的bbox
def postprocess(prediction, num_classes=2, conf_thre=0.3, nms_thre=0.45):
box_corner = prediction.new(prediction.shape)
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
prediction[:, :, :4] = box_corner[:, :, :4]
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction):
# If none are remaining => process next image
if not image_pred.size(0):
continue
# Get score and class with highest confidence
class_conf, class_pred = torch.max(
image_pred[:, 5:5 + num_classes], 1, keepdim=True)
conf_mask = (image_pred[:, 4] * class_conf.squeeze() >= conf_thre).squeeze()
# Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred)
detections = torch.cat(
(image_pred[:, :5], class_conf, class_pred.float()), 1)
detections = detections[conf_mask]
if not detections.size(0):
continue
# Iterate through all predicted classes
unique_labels = detections[:, -1].unique()
for c in unique_labels:
# Get the detections with the particular class
detections_class = detections[detections[:, -1] == c]
nms_out_index = torchvision.ops.nms(
detections_class[:, :4], detections_class[:, 4]*detections_class[:, 5], nms_thre)
detections_class = detections_class[nms_out_index]
if output[i] is None:
output[i] = detections_class
else:
output[i] = torch.cat((output[i], detections_class))
return output
outputs = postprocess(outputs, 2, 0.01, 0.35)
outputs = outputs[0].cpu().data
bboxes = outputs[:, 0:4]
bboxes[:, 0::2] *= width / test_size[0]
bboxes[:, 1::2] *= height / test_size[1]
bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
cls = outputs[:, 6]
scores = outputs[:, 4] * outputs[:, 5]最后我们将处理得到的结果可视化出来
def vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None, color=None):
colors = torch.FloatTensor([[1,0,1],[0,0,1],[0,1,1],[0,1,0],[1,1,0],[1,0,0]]);
def get_color(c, x, max_val):
ratio = float(x)/max_val * 5
i = int(math.floor(ratio))
j = int(math.ceil(ratio))
ratio = ratio - i
r = (1-ratio) * colors[i][c] + ratio*colors[j][c]
return int(r*255)
width = img.shape[1]
height = img.shape[0]
for i in range(len(boxes)):
box = boxes[i]
cls_conf = scores[i]
if cls_conf < conf:
continue
x1 = int(box[0])
y1 = int(box[1])
x2 = int(box[0]+box[2])
y2 = int(box[1]+box[3])
if color:
rgb = color
else:
rgb = (255, 0, 0)
if class_names is not None:
cls_conf = scores[i]
cls_id = int(cls_ids[i])
class_name = class_names[cls_id]
classes = len(class_names)
offset = cls_id * 123456 % classes
red = get_color(2, offset, classes)
green = get_color(1, offset, classes)
blue = get_color(0, offset, classes)
if color is None:
rgb = (red, green, blue)
img = cv2.putText(img, '%s: %.2f'%(class_name,cls_conf), (x1,y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, rgb, 2)
img = cv2.rectangle(img, (x1,y1), (x2,y2), rgb, 1)
return img
pred_im = vis(ori_im, bboxes.numpy(), scores.numpy(), cls.numpy(), conf=0.3, class_names=VOC_CLASSES)
plt.rcParams['figure.figsize'] = (20, 12)
plt.imshow(pred_im[:,:,::-1])<matplotlib.image.AxesImage at 0x7f0c4cfc7950>
0015e6a27a5a1671b03db56e2718f19需要说明的是本教程只是用于入门学习，输入图像尺寸选择416x416，没有加入复杂的数据增广算法、网络训练优化tricks，这些因素都导致模型存在人群密集情况漏检以及小人头漏检问题。如果希望获得更好的效果，可以使用人脸检测模型，混入这个口罩数据集，这样能够利用更多的人脸数据。最后对检测出的人脸图像做2分类即可。