MTCNN人脸检测 pytorch实现 全流程详解
概述
MTCNN，Multi-task convolutional neural network（多任务卷积神经网络）。该模型利用级联思想，化繁为简，通过三个级联的网络：P-Net（Proposal Network）、R-Net（Refine Network）和O-Net（Output Network），逐层筛选。在每个网络都非常简单，易于训练的同时，实现了高精度的人脸检测。
思路
人脸检测属于单类多目标检测，相对于单类单目标，实现中有以下问题：
- 人脸数目
因为不知道1张图中需要检测多少人脸，也就不能通过网络直接输出预测值。
解决方法：像卷积操作中卷积核的滑动一样，以一定的大小和步长扫描整张图片，将扫描到的小图输入网络做检测。 - 人脸大小
紧接着的问题是：如何确定“卷积核”大小？
因为需要检测的人脸大小也不确定，所以考虑应该用不同大小的“卷积核”去扫描。
反之，也可以保持“卷积核”不变，通过数次缩小原图（图像金字塔）来实现。 - 精度不高
如果像普通单类单目标任务一样，直接回归真实人脸框，那么检测中一旦错过了合适的框，就错过该目标了。
因此MTCNN中一个精妙的点在于：不直接学习真实框，而是学习偏移框，由偏移量反算回真实框。
这样做的好处是：在人脸目标周围，会检测到多个偏移框，从而反算回多个预测框，大大增加了容错率。 - 一大堆框
这样就可以将不同大小的扫描图片传入P-Net检测，在一大堆重叠的预测框中，我们需要的，只应该是其中最确定的那一个，这就依赖于NMS操作了 。
数据处理
选择CelebA人脸数据集，标签为人脸框左上角坐标：x1、y1，宽高：w、h，及5个人脸特征点坐标。
- 对真实框中心坐标做随机偏移，最大边长+随机量作为新边长，生成正方形框。可以选择将1张图片增样成数张。
- 将正方形框和原真实框做 IOU，通过设定 IOU 的阈值，将生成框划分为正样本、负样本、部分样本。
- 裁下正方形框，resize成尺寸12、24、48，分别用于P、R、O网络训练。
- 计算生成框坐标、5个人脸特征点坐标的偏移量，负样本偏移量均设为0，因其不参与坐标回归训练。
计算偏移量作用相当于归一化，加快收敛。 - 将置信度、2个生成框坐标偏移量、5个人脸特征点坐标偏移量，共15个值写入txt文件作为训练标签。
生成框坐标偏移量计算如下图，其余特征点坐标偏移量计算方法同 (_x1，_y1)，对左上角点做偏移即可，也可以自己找一个点做偏移，如右下角点或矩形中点，测试时按对应点反算坐标。
我处理后的3种样本数据集，每种尺寸20万张，共60万张。
正样本：
部分样本：
负样本：
工具
非极大值抑制
交并比（Intersection-over-Union，IOU），两个图形交集与并集的比值，表示重叠程度。
需要注意的是：在最后的O网络输出时，将把 IOU 计算公式中分母改为A和B中的面积较小值，从而去除候选框大框套小框的情况。
非极大值抑制（Non-Maximum Suppression，NMS）
作用：在同一个人脸的一堆候选框中，选出最可信的一个。
步骤1：将候选框按置信度降序排序
步骤2：第一个候选框依次与之后的框做 IOU， 大于设定阈值的框（认为同一个人脸）舍弃
步骤3：保留第一个候选框，剩余的候选框重复步骤1~步骤3，直至剩下一个框，保留。
最后保留下来的候选框就是 NMS 处理后的结果。
Soft-NMS 是对 NMS 的一种改进。
伪代码如图：
B: 初始检测框集合，S：对应检测框的分数， Nt ：IOU的阈值，M ：得分最高的检测框
NMS 的处理是将大于 IOU 阈值的框直接舍弃（分数置0），这样容易错过一些重合度较大的框。
Soft-NMS 的思路：根据 IOU 的值降低框的得分，最后根据分数阈值统一删除。
降低置信度的方法有两种：
- 线性加权（不连续）
- 高斯加权（连续）
这两种方法在实验中效果差别不大，因此代码中我选择第一种更简洁的形式。
def nms(boxes, thresh=0.3, is_min=False, softnms=False):
if boxes.shape[0] == 0:
return np.array([])
_boxes = boxes[(-boxes[:, 14]).argsort()] # 按置信度排序
r_boxes = []
while _boxes.shape[0] > 1:
a_box = _boxes[0]
b_boxes = _boxes[1:]
score = b_boxes[:, 14]
r_boxes.append(a_box)
if softnms:
score_thresh = 0.5
# IOU>阈值的框 置信度衰减
t_idx = np.where(iou(a_box, b_boxes, is_min) > thresh)
score[t_idx] *= (1 - iou(a_box, b_boxes, is_min))[t_idx]
# 删除分数<阈值的框
_boxes = np.delete(b_boxes, np.where(score < score_thresh), axis=0)
else:
# 筛选IOU<阈值的框
index = np.where(iou(a_box, b_boxes, is_min) < thresh)
_boxes = b_boxes[index]
# 剩余最后1个框 保留
if _boxes.shape[0] > 0:
r_boxes.append(_boxes[0])
# 把list组装成矩阵
return np.stack(r_boxes)
下图为 NMS 和Soft-NMS 效果对比
图像金字塔
保持检测框尺寸不变（如12）， 将待检测图像按一定比例连续缩小，直到图像短边尺寸小于检测框尺寸。
每次缩小的图片上再用检测框滑动检测，这样就能涵盖原图上大于检测框尺寸的所有人脸。
缩小比例一般用0.6或0.7，官方推荐0.709，是根据让每次图像面积缩小一半得出的。
网络结构
P网络采用全卷积结构，最后一层用1×1 的卷积层代替全连接层，这是因为P网络训练时使用的是12×12的图片，但是待检测的图像尺寸并不是，如果最后一层用全连接，shape转换会报错。
P网络输出的建议框，尺寸resize成24、48之后，才传入R、O网络，所以R、O网络最后一层用全连接即可。
我的网络增加了BatchNormal层，因此卷积中的偏置参数可以去掉。
其中池化层还可以用卷积层替代，网络表达能力更好。
import torch
import torch.nn as nn
import torchsummary
class PNet(nn.Module):
def __init__(self):
super(PNet, self).__init__()
self.pre_layer = nn.Sequential(
nn.Conv2d(3, 10, kernel_size=3, stride=1, bias=False), # 10*10*10
nn.BatchNorm2d(10),
nn.PReLU(),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0), # 5*5*10
nn.Conv2d(10, 16, kernel_size=3, stride=1, bias=False), # 3*3*16
nn.BatchNorm2d(16),
nn.PReLU(),
nn.Conv2d(16, 32, kernel_size=3, stride=1, bias=False), # 1*1*32
nn.BatchNorm2d(32),
nn.PReLU()
)
self.conv_1 = nn.Conv2d(32, 1, kernel_size=1, stride=1)
self.conv_2 = nn.Conv2d(32, 14, kernel_size=1, stride=1)
def forward(self, x):
y = self.pre_layer(x)
cls = torch.sigmoid(self.conv_1(y))
offset = self.conv_2(y)
return cls, offset
class RNet(nn.Module):
def __init__(self):
super(RNet, self).__init__()
self.pre_layer = nn.Sequential(
nn.Conv2d(3, 28, kernel_size=3, stride=1, bias=False), # 22*22*28
nn.BatchNorm2d(28),
nn.PReLU(),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1), # 11*11*28
nn.Conv2d(28, 48, kernel_size=3, stride=1, bias=False), # 9*9*48
nn.BatchNorm2d(48),
nn.PReLU(),
nn.MaxPool2d(kernel_size=3, stride=2), # 4*4*48
nn.Conv2d(48, 64, kernel_size=2, stride=1, bias=False), # 3*3*64
nn.BatchNorm2d(64),
nn.PReLU()
)
self.fc1 = nn.Sequential(
nn.Linear(64 * 3 * 3, 128),
nn.PReLU()
)
self.fc2_1 = nn.Linear(128, 1)
self.fc2_2 = nn.Linear(128, 14)
def forward(self, x):
y = self.pre_layer(x)
y = y.view(y.size(0), -1)
y = self.fc1(y)
cls = torch.sigmoid(self.fc2_1(y))
offset = self.fc2_2(y)
return cls, offset
class ONet(nn.Module):
def __init__(self):
super(ONet, self).__init__()
self.pre_layer = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, stride=1, bias=False), # 46*46*32
nn.BatchNorm2d(32),
nn.PReLU(),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1), # 23*23*32
nn.Conv2d(32, 64, kernel_size=3, stride=1, bias=False), # 21*21*64
nn.BatchNorm2d(64),
nn.PReLU(),
nn.MaxPool2d(kernel_size=3, stride=2), # 10*10*64
nn.Conv2d(64, 64, kernel_size=3, stride=1, bias=False), # 8*8*64
nn.BatchNorm2d(64),
nn.PReLU(),
nn.MaxPool2d(kernel_size=2, stride=2), # 4*4*64
nn.Conv2d(64, 128, kernel_size=2, stride=1, bias=False), # 3*3*128
nn.BatchNorm2d(128),
nn.PReLU()
)
self.fc1 = nn.Sequential(
nn.Linear(128 * 3 * 3, 256),
nn.PReLU()
)
self.fc2_1 = nn.Linear(256, 1)
self.fc2_2 = nn.Linear(256, 14)
def forward(self, x):
y = self.pre_layer(x)
y = y.view(y.size(0), -1)
y = self.fc1(y)
cls = torch.sigmoid(self.fc2_1(y))
offset = self.fc2_2(y)
return cls, offset
if __name__ == '__main__':
# torchsummary.summary(PNet().cuda(), (3, 12, 12))
torchsummary.summary(RNet().cuda(), (3, 24, 24))
# torchsummary.summary(ONet().cuda(), (3, 48, 48))
训练流程
用生成的12、24、48尺寸数据分别训练P网络、R网络、O网络，3个网络能力递增，可并行训练。
损失 Loss 分为两部分：
- 置信度损失
置信度即判断是否为人脸，由正样本和负样本训练，损失函数选择二值交叉熵损失(BCELoss)。 - 偏移量损失
偏移量用来回归人脸框和特征点坐标，由正样本和部分样本训练，损失函数选择均方差损失(MSELoss)。
置信度损失与偏移量损失的加权和，即为总损失。
P、R网络更加注重分类精度训练，找到可靠的候选区域，回归精度在O网络中达到要求即可。
权重 alpha 值可自行尝试，我选择的值分别为0.8、0.7、0.5。
每5轮绘制训练损失、分类准确率、偏移量r2分数的图像。
import os
import torch
import torch.optim as optim
import numpy as np
from torch.utils.data import DataLoader
from torch import nn
from sample import FaceDataset
from sklearn.metrics import r2_score
import matplotlib.pyplot as plt
class Trainer:
def __init__(self, net, save_path, dataset_path):、
self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
self.net = net.to(self.device)
self.save_path = save_path
self.dataset_path = dataset_path
self.cls_loss_fn = nn.BCELoss()
self.offset_loss_fn = nn.MSELoss()
self.optim_adam = optim.Adam(self.net.parameters())
self.optim_sgd = optim.SGD(self.net.parameters(), lr=5e-4, momentum=0.9)
if os.path.exists(self.save_path):
# load_state_dict 接收字典对象
net.load_state_dict(torch.load(self.save_path))
else:
print("No Param")
def train(self, epochs=int(), alpha=0.5):
dataset = FaceDataset(self.dataset_path)
loader = DataLoader(dataset, batch_size=2048, shuffle=True, num_workers=4, pin_memory=True)
loss_list = []
acc_list = []
r2_list = []
for epoch in range(epochs):
true_num = 0
all_num = 0
train_loss = 0
r2_arr = np.array([])
optimizer = self.optim_adam if epoch < (epochs / 2) else self.optim_sgd
print("epoch:\t{}".format(epoch + 1))
for i, (data, cls, offset) in enumerate(loader):
data, cls, offset = data.to(self.device), cls.to(self.device), offset.to(self.device)
output_cls, output_offset = self.net(data)
# P输出格式 NCHW (N,1,1,1) R/O输出格式 NC (N,1)
output_cls = output_cls.view(-1, 1)
output_offset = output_offset.view(-1, 14)
# 分类损失
category = torch.masked_select(cls, cls < 2)
output_cls = torch.masked_select(output_cls, cls < 2)
pred = torch.where(output_cls > 0.5, torch.tensor(1).cuda(), torch.tensor(0).cuda())
true_num += (pred == category).sum().item()
all_num += category.shape[0]
cls_loss = self.cls_loss_fn(output_cls, category)
# 偏移量损失
offset = torch.masked_select(offset, cls > 0)
output_offset = torch.masked_select(output_offset, cls > 0)
offset_loss = self.offset_loss_fn(output_offset, offset)
loss = alpha * cls_loss + (1 - alpha) * offset_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
cls_loss = cls_loss.cpu().item()
offset_loss = offset_loss.cpu().item()
loss = loss.cpu().item()
train_loss += (loss * cls.shape[0])
if i % 20 == 0:
r2_ = r2_score(offset.detach().cpu().numpy(), output_offset.detach().cpu().numpy())
r2_arr = np.append(r2_arr, r2_)
print("loss: {:.8f}\tcls_loss: {:.8f}\toffset_loss: {:.8f}".format(loss, cls_loss, offset_loss))
train_loss /= len(dataset)
loss_list.append(train_loss)
print("loss:\t\t{:.6f}".format(train_loss))
accuracy = true_num / all_num
acc_list.append(accuracy)
print("accuracy:\t{:.4f}%".format(accuracy * 100))
r2 = r2_arr.mean()
r2_list.append(r2)
print("r2_score:\t{:.6f}".format(r2))
torch.save(self.net.state_dict(), self.save_path)
print("Saved successfully\n")
if (epoch + 1) % 5 == 0:
plt.figure(figsize=(15, 5))
plt.subplot(131)
plt.plot(loss_list)
plt.title('loss')
plt.subplot(132)
plt.plot(acc_list)
plt.title('accuracy')
plt.subplot(133)
plt.plot(r2_list)
plt.title('r2_score')
plt.savefig('graph/loss{}.png'.format(epoch + 1))
检测流程
检测流程是难点，具体如下：
坐标反算
坐标反算分为两部分：
- P网络中由输出特征图反算回偏移框坐标。
- 偏移框根据偏移量反算回原图坐标
先看第一部分，为了方便画图，暂时忽略5个特征点坐标。P网络输出的是全卷积后的特征图，分为两部分：置信度 cls (1,1,H,W) 和 偏移量 offset (1,4,H,W)，共5个通道。
P网络作为全卷积网络，对输入尺寸大于12×12的图片，相当于用12×12的卷积核做卷积，这个过程就巧妙地实现了用检测框扫描全图。需要注意的是，这个卷积过程的步长，相当于P网络中每层卷积的步长乘积，即等于2。
如下图，假设测试图片尺寸为14×16，输入P网络后，输出特征图shape为 (1, 5, 2, 3) 。特征图中红点 (H, W) 的坐标为 (1, 2)，对应卷积前原图中的红框区域，我们的任务就是由 (H, W) 的坐标 (1, 2) 反算回红框在原图中的坐标。
左上角坐标 = (特征图坐标 × 卷积步长)
右下角坐标 = (左上角坐标 + 卷积核大小)
在该栗子中，左上角坐标 = (2, 1) × 2 = (4, 2)，右下角坐标 = (4, 2) + (12, 12) = (16, 14) 。
注意到 (H, W) 应该对应坐标 (y, x)，所以计算时候需要交换位置。
因为原图是经过图像金字塔才输入网络的，所以还需除以相对于原图的缩小比例scale，才能得到原图中的偏移框坐标。
第二部分比较简单，根据生成数据时偏移量的计算公式，由偏移框坐标 (_x1, _y1, _x2, _y2) 、偏移框尺寸_side_len、网络输出的偏移量offset，反算即可。
对P网络得到的预测框，根据置信度进行筛选。
因为P网络的精度最低，任务是要选出所有可能的建议框传给后边的网络，所以筛选较为宽松。比如阈值可以给0.6，所有预测分类置信度大于0.6的框就都留下。
剩余的框进行坐标反算，然后经过 NMS 处理，去除重合度过大的框。这里的 IOU 阈值也是自己给定。
最后将框按长边填充成正方形，resize成24×24，传入R网络。
R网络的处理就较简单了，将输入框重新计算置信度和偏移量，同样进行置信度筛选、坐标反算、NMS、填充正方形、resize成48×48，再传入O网络。
O网络处理和R网络几乎一样，再次提高检测精度，坐标反算后在原图中画出预测框。