pytorch项目实战之实时人脸属性检测系统
简介
本项目采用CelebA人脸属性数据集训练人脸属性分类模型，使用mediapipe进行人脸检测，使用onnxruntime进行模型的推理，最终在intel的奔腾cpu上实现30-100帧完整的实时人脸属性识别系统。
ps:本来是打算写成付费专栏的，毕竟这是个只要稍微修改就可以商业化的系统，不是那些玩具例子，但考虑到阅读量马上破十万了，就作为纪念作发出来吧。
python包环境
训练环境：
python 3.9.12
torch 1.12.0+cu116
torchvision 0.13.0+cu116
导出模型相关：
onnx 1.12.0
onnxruntime 1.14.0
部署模型环境：
onnxruntime 1.14.0
cv2(opencv-python) 4.7.0
mediapipe 0.9.0.1
python3.10.6
注：版本没有特定要求，写出来只是方便查bug，一般拥有以上的包即可完成本项目，没有包请自行pip安装。
数据集准备
手动下载CelebA数据集
输入截图上面这个网址（不直接放网址是因为外链会被识别为低质量文章的奇葩规定），然后点击图上baidu drive，输入提取码后：
把以上文件夹下的每一个东西都下载下来放到一起，比如我这边下载完后构造如下图所示，一个都不能少和多：
img_align_celeba文件夹是压缩包解压完后的，下面直接就是人脸的图片，没有次级目录：
装载数据集
首先使用torchvision直接加载数据集，划分训练集，测试集。这里download设为False表示你已经手动下载好了，不需要自动下载（速度很慢）。此外root路径为数据集的根路径，我这里所有东西是放在D:/face/celeba下，那么根路径就是写D:/face。
from torchvision import datasets
import torchvision.transforms as transforms
train_dataset = datasets.CelebA(root="D:/face",
split='train',
transform = transforms.Compose([
transforms.CenterCrop(128),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])]),
download=False)
test_dataset = datasets.CelebA(root="D:/face",
split='test',
transform = transforms.Compose([
transforms.CenterCrop(128),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])]),
download=False)
这里对数据集进行了三个预处理，因为涉及到后面推理环节摆脱torch进行相同的处理所以需要详细解释：
| 语句 | 意义 |
|---|---|
| transforms.CenterCrop(128) | 中心裁剪128x128图片 |
| transforms.ToTensor() | 把0-255的像素值除以255转为0-1 |
| transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5, 0.5, 0.5])]) | 对三个通道减去0.5平均值并除以0.5标准差 |
然后我们使用torch的dataloader类来把数据集正式装入可供读取：
import torch
train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, shuffle=True, batch_size=128,num_workers=4)
test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, shuffle=False, batch_size=128,num_workers=4)
batch_size我这里调了128，num_workers调了4，这两个指标是影响训练速度的，如果电脑配置不行请适当调小，配置很行请调大加速训练，数据量本身还是很大的，一次epoch大概2分钟。
定义模型类
模型采用2017年的slimnet架构，是一个非常轻量的网络，对slimnet感兴趣的请自行搜索相关论文，这里直接放代码了：
import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvBNReLU(nn.Sequential):
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
padding = (kernel_size - 1) // 2
super(ConvBNReLU, self).__init__(
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
nn.BatchNorm2d(out_planes),
nn.ReLU(inplace=True)
)
class DWSeparableConv(nn.Module):
def __init__(self, inp, oup):
super().__init__()
self.dwc = ConvBNReLU(inp, inp, kernel_size=3, groups=inp)
self.pwc = ConvBNReLU(inp, oup, kernel_size=1)
def forward(self, x):
x = self.dwc(x)
x = self.pwc(x)
return x
class SSEBlock(nn.Module):
def __init__(self, inp, oup):
super().__init__()
out_channel = oup * 4
self.pwc1 = ConvBNReLU(inp, oup, kernel_size=1)
self.pwc2 = ConvBNReLU(oup, out_channel, kernel_size=1)
self.dwc = DWSeparableConv(oup, out_channel)
def forward(self, x):
x = self.pwc1(x)
out1 = self.pwc2(x)
out2 = self.dwc(x)
return torch.cat((out1, out2), 1)
class SlimModule(nn.Module):
def __init__(self, inp, oup):
super().__init__()
hidden_dim = oup * 4
out_channel = oup * 3
self.sse1 = SSEBlock(inp, oup)
self.sse2 = SSEBlock(hidden_dim * 2, oup)
self.dwc = DWSeparableConv(hidden_dim * 2, out_channel)
self.conv = ConvBNReLU(inp, hidden_dim * 2, kernel_size=1)
def forward(self, x):
out = self.sse1(x)
out += self.conv(x)
out = self.sse2(out)
out = self.dwc(out)
return out
class SlimNet(nn.Module):
def __init__(self, num_classes):
super().__init__()
self.conv = ConvBNReLU(3, 96, kernel_size=7, stride=2)
self.max_pool0 = nn.MaxPool2d(kernel_size=3, stride=2)
self.module1 = SlimModule(96, 16)
self.module2 = SlimModule(48, 32)
self.module3 = SlimModule(96, 48)
self.module4 = SlimModule(144, 64)
self.max_pool1 = nn.MaxPool2d(kernel_size=3, stride=2)
self.max_pool2 = nn.MaxPool2d(kernel_size=3, stride=2)
self.max_pool3 = nn.MaxPool2d(kernel_size=3, stride=2)
self.max_pool4 = nn.MaxPool2d(kernel_size=3, stride=2)
self.gap = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(192, num_classes)
def forward(self, x):
x = self.max_pool0(self.conv(x))
x = self.max_pool1(self.module1(x))
x = self.max_pool2(self.module2(x))
x = self.max_pool3(self.module3(x))
x = self.max_pool4(self.module4(x))
x = self.gap(x)
x = torch.flatten(x, 1)
x = self.fc(x)
return x
device = torch.device('cuda')
model = SlimNet(num_classes=40).to(device=device)
device这里我用了cuda训练，如果你没有英伟达显卡请写cpu，num_classes=40表示最后输出有40个人脸属性特征，以下是我翻译过后的特征，方便理解：
可以看到我们最后是个多分类任务，但这里网络输出的不是0,1的分类，也不是0-1的概率值，需要在预测时先用sigmoid函数转换到0-1的可能性，再以0.5分类。
训练网络
中间的训练过程就不细讲了，都是pytorch的八股文，需要注意的是选取损失函数时要明白它的数据类型，这里target是整数型，需要转为double型才能和score计算loss。本次训练手动设定了一个best_acc表示最好的准确度，一旦在训练中对测试集评估时发现准确度比它更好，就会自动保存当前模型。
loss_criterion = nn.BCEWithLogitsLoss() #定义损失函数
optimizer = torch.optim.Adam(model.parameters(), lr = 0.001) #定义优化器
best_acc = 0.90325 #最好的在测试集上的准确度，可手动修改
seed = 18203861252700 #固定起始种子
for epoch in range(50): #训练五十轮
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True#以上这些都是企图固定种子，但经过测试只能固定起始种子，可删掉
total_train = 0 #总共的训练图片数量，用来计算准确率
correct_train = 0 #模型分类对的训练图片
running_loss = 0 #训练集上的loss
running_test_loss = 0 #测试集上的loss
total_test = 0 #测试的图片总数
correct_test = 0 #分类对的测试图片数
model.train() #训练模式
for data, target in train_dataloader:
data = data.to(device=device)
target = target.type(torch.DoubleTensor).to(device=device)
score = model(data)
loss = loss_criterion(score, target)
running_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
sigmoid_logits = torch.sigmoid(score)
predictions = sigmoid_logits > 0.5 #使结果变为true,false的数组
total_train += target.size(0) * target.size(1)
correct_train += (target.type(predictions.type()) == predictions).sum().item()
model.eval() #测试模式
with torch.no_grad():
for batch_idx, (images,labels) in enumerate(test_dataloader):
images, labels = images.to(device), labels.type(torch.DoubleTensor).to(device)
logits = model.forward(images)
test_loss = loss_criterion(logits, labels)
running_test_loss += test_loss.item()
sigmoid_logits = torch.sigmoid(logits)
predictions = sigmoid_logits > 0.5
total_test += labels.size(0) * labels.size(1)
correct_test += (labels.int() == predictions.int()).sum().item()
test_acc = correct_test/total_test
if test_acc > best_acc:
best_acc = test_acc
torch.save(model,f"model_{test_acc*100}.pt")
print(f"For epoch : {epoch} training loss: {running_loss/len(train_dataloader)}")
print(f'train accruacy is {correct_train*100/total_train}%')
print(f"For epoch : {epoch} test loss: {running_test_loss/len(test_dataloader)}")
print(f'test accruacy is {test_acc*100}%')
训练结束后可以在目录下找到以下文件，就是我们训练过程中发现的好模型，用来后续导出：
模型导出onnx
这里最好接着上面在同一个notebook里写，如果另外保存的话，需要把上面定义网络的部分复制粘贴进去，不然会找不到网络的定义。
torch_model = torch.load("model_90.56845506462278.pt", map_location='cpu')
torch_model.eval()
x = torch.randn(1, 3, 128, 128, requires_grad=True) #随机128x128输入
torch_out = torch_model(x)
print(torch_out)
# 导出模型
torch.onnx.export(torch_model, # 需要导出的模型
x, # 模型输入
"cpu.onnx", # 保存模型位置
export_params=True, # 保存训练参数
opset_version=10, # onnx的opset版本
do_constant_folding=True, # 是否进行常量折叠优化，这里开关都一样
input_names = ['input'], # 输入名字
output_names = ['output'], # 输出名字
)
ort_session = onnxruntime.InferenceSession("cpu.onnx",providers=['CPUExecutionProvider'])
#尝试进行推理看是否报错
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)
print(ort_outs[0])
# 比较onnx模型推理的结果和torch推理的结果误差是否在可容忍范围内
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
print("Exported model has been tested with ONNXRuntime, and the result looks good!")
一切没问题的话会在目录下找到cpu.onnx文件：
图片最小推理示例
我们直接把数据集中的第一张图片复制过来并重名为test_face.jpg，
进行推理：
from PIL import Image
import torchvision.transforms as transforms
import onnxruntime
import torch
import numpy as np
img = Image.open("test_face.jpg")
comp = transforms.Compose([transforms.CenterCrop(128),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),]) #torch的预处理
img = comp(img)
img.unsqueeze_(0)
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
x = to_numpy(img)
ort_session = onnxruntime.InferenceSession("cpu.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: x}
ort_outs = ort_session.run(None, ort_inputs)
def sigmoid_array(x): #使用sigmoid转换概率值
return 1 / (1 + np.exp(-x))
result = sigmoid_array(ort_outs[0]) > 0.5
list_attr_cn = np.array(["早上刚刮下午长出来的一点胡子","拱形眉毛","有吸引力的","眼袋","秃头","刘海","大嘴唇"
,"大鼻子","黑发","金发","模糊的","棕发","浓眉","圆胖","双下巴","眼镜","山羊胡子","灰白发","浓妆","高高的颧骨",
"男性","嘴微微张开","胡子","眯眯眼","没有胡子","鹅蛋脸","苍白皮肤","尖鼻子","后退的发际线","红润脸颊",
"鬓角","微笑","直发","卷发","耳环","帽子","口红","项链","领带","年轻"])
print(list_attr_cn[result[0]])
结果如下：
[‘拱形眉毛’ ‘有吸引力的’ ‘棕发’ ‘浓妆’ ‘高高的颧骨’ ‘嘴微微张开’ ‘没有胡子’ ‘尖鼻子’ ‘微笑’ ‘口红’ ‘年轻’]
有了玩具例子，只要再加上亿点点细节，就可以完成整个系统了。
视频实时人脸检测
直接放代码，如果要细细解释得讲很多，不如直接代码注释看起来方便，总之就是用opencv来读取视频并画文字，使用Mediapipe来进行快速人脸识别。
import onnxruntime
import time
import numpy as np
import cv2
import mediapipe as mp
mp_face_detection = mp.solutions.face_detection
mp_drawing = mp.solutions.drawing_utils
cap = cv2.VideoCapture("test_face3.mp4") #视频输入，如果需要摄像头，请改成数字0，并修改下面的break为continue
ort_session = onnxruntime.InferenceSession("cpu.onnx")
list_attr_en = np.array(["5_o_Clock_Shadow","Arched_Eyebrows","Attractive","Bags_Under_Eyes","Bald",
"Bangs","Big_Lips","Big_Nose","Black_Hair","Blond_Hair","Blurry","Brown_Hair",
"Bushy_Eyebrows","Chubby","Double_Chin","Eyeglasses","Goatee","Gray_Hair",
"Heavy_Makeup","High_Cheekbones","Male","Mouth_Slightly_Open","Mustache","Narrow_Eyes",
"No_Beard","Oval_Face","Pale_Skin","Pointy_Nose","Receding_Hairline","Rosy_Cheeks",
"Sideburns","Smiling","Straight_Hair","Wavy_Hair","Wearing_Earrings","Wearing_Hat",
"Wearing_Lipstick","Wearing_Necklace","Wearing_Necktie","Young"]) #英文原版属性
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) #获取视频宽度
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) #获取视频高度
fps = cap.get(cv2.CAP_PROP_FPS) #获取视频FPS，如果是实时摄像头请手动设定帧数
out = cv2.VideoWriter('output3.avi', cv2.VideoWriter_fourcc(*"MJPG"), fps, (width,height)) #保存视频，没需求可去掉
def cv2_preprocess(img): #numpy预处理和torch处理一样
img = cv2.resize(img, (128, 128), interpolation=cv2.INTER_NEAREST)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
mean = [0.5,0.5,0.5] #一定要是3个通道，不能直接减0.5
std = [0.5,0.5,0.5]
img = ((img / 255.0 - mean) / std)
img = img.transpose((2,0,1)) #hwc变为chw
img = np.expand_dims(img, axis=0) #3维到4维
img = np.ascontiguousarray(img, dtype=np.float32) #转换浮点型
return img
def sigmoid_array(x): #sigmoid函数手动设定
return 1 / (1 + np.exp(-x))
def result_inference(input_array): #推理环节
ort_inputs = {ort_session.get_inputs()[0].name: input_array}
ort_outs = ort_session.run(None, ort_inputs)
possibility = sigmoid_array(ort_outs[0]) > 0.5
result = list_attr_en[possibility[0]]
return result
with mp_face_detection.FaceDetection(
model_selection=1, min_detection_confidence=0.5) as face_detection:
#人脸识别，1为通用模型，0为近距离模型
while cap.isOpened():
a1 = time.time()
success, image = cap.read()
if not success:
print("Ignoring empty camera frame.")
break
image.flags.writeable = False #据说这样写可以加速人脸识别推理
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
results = face_detection.process(image)
image.flags.writeable = True
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
image2 = image.copy() #copy复制，因为cv2会直接覆盖原有数组
if results.detections:
for detection in results.detections:
mp_drawing.draw_detection(image, detection)
image_rows, image_cols, _ = image.shape
location = detection.location_data.relative_bounding_box #获取人脸位置
start_point = mp_drawing._normalized_to_pixel_coordinates(location.xmin, location.ymin,image_cols,image_rows) #获取人脸左上角的点
end_point = mp_drawing._normalized_to_pixel_coordinates(location.xmin + location.width, location.ymin + location.height,image_cols,image_rows) #获取右下角的点
x1,y1 = start_point #左上点坐标
x2,y2 = end_point #右下点坐标
img_infer = image2[y1-70:y2,x1-50:x2+50].copy() #为了营造相似环境，把左上角和右上角的点连线囊括的区域扩大提高准确度
img_infer = cv2_preprocess(img_infer)
result = result_inference(img_infer)
# # cv2.imshow('test',img_infer)
# if cv2.waitKey(5) & 0xFF == 27:
# break
for i in range(0,len(result)):
image = cv2.putText(image, result[i],(x1,y1+i*40), cv2.FONT_HERSHEY_SIMPLEX,
1, (255,255,255), 1, cv2.LINE_AA) #画文字，一行一行画，如果想要中文请自行编译安装freetype版opencv，不推荐用别的库包裹转换中文，速度慢
a2 = time.time()
out.write(image)
print(f'one pic time is {a2 - a1} s')
此外，mediapipe本身提供的功能是自带画人脸关键点的，如果不想要关键点可以和我一样把它注释掉：
最终成果
推理时间测试了多个视频，一帧完整的从读取到检测画图画完的时间大约在0.01-0.035之间，换算下来每秒可以达到30-100帧。
从下图可以看到识别出来的属性还是挺准的，图上是一位棕色头发没有胡子的有吸引力的年轻女性。属性条有变化是因为人脸识别捕捉的区域不同，也没有加上人脸跟踪的算法。