facenet-pytorch库的简单使用

最近尝试facenet做识别，没有从头复现，刚好在GitHub找到一个可以已经封装好的repo，使用起来也特别方便。项目地址 https://github.com/timesler/facenet-pytorch 。安装使用它只需直接pip install facenet-pytorch即可。

facenet-pytorch库使用mtcnn进行人脸检测和InceptionResnetV1进行图像到欧式空间的向量映射。进行使用之前需要建立要识别的人脸特征向量数据库。结合timesler提供的例子，总结了人脸数据库制作的代码如下，提前将数据库中的人脸向量与其名字保存为本地文件，这样对新图识别的时候直接加载这个文件即可。

该库使用的网络预训练权重在readme里面都有。

第一步：制作自己的数据库。

# 制作人脸特征向量的数据库 最后会保存两个文件，分别是数据库中的人脸特征向量和对应的名字。当然也可以保存在一起
from facenet_pytorch import MTCNN, InceptionResnetV1
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np
import pandas as pd
import os
workers = 0 if os.name == 'nt' else 4
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))
mtcnn = MTCNN(
    image_size=160, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
    device=device
)
# InceptionResnetV1提供了两个预训练模型，分别在vggface数据集和casia数据集上训练的。
# 预训练模型如果不手动下载，可能速度会很慢，可以从作者给的谷歌云链接下载，然后放到C:\Users\你的用户名\.cache\torch\checkpoints这个文件夹下面
# 如果是linux系统，那么存放在/home/你的用户名/.cache/torch/checkpoints下面
resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)

def collate_fn(x):
    return x[0]
# 将所有的单人照图片放在各自的文件夹中，文件夹名字就是人的名字,存放格式如下
'''
--orgin
  |--zhangsan
     |--1.jpg
     |--2.jpg
  |--lisi
     |--1.jpg
     |--2.jpg
'''
dataset = datasets.ImageFolder('./database/orgin')  #加载数据库 
dataset.idx_to_class = {i:c for c, i in dataset.class_to_idx.items()}
loader = DataLoader(dataset, collate_fn=collate_fn, num_workers=workers)
aligned = []  # aligned就是从图像上抠出的人脸，大小是之前定义的image_size=160
names = []
i= 1
for x, y in loader:
    path = './database/aligned/{}/'.format(dataset.idx_to_class[y])  # 这个是要保存的人脸路径
    if not os.path.exists(path):
        i = 1
        os.mkdir(path)
    # 如果要保存识别到的人脸，在save_path参数指明保存路径即可,不保存可以用None
    x_aligned, prob = mtcnn(x, return_prob=True,save_path= path+ '/{}.jpg'.format(i))
    i = i+1
    if x_aligned is not None:
        print('Face detected with probability: {:8f}'.format(prob))
        aligned.append(x_aligned) 
        names.append(dataset.idx_to_class[y])

aligned = torch.stack(aligned).to(device)
embeddings = resnet(aligned).detach().cpu()   # 提取所有人脸的特征向量，每个向量的长度是512
# 两两之间计算混淆矩阵
dists = [[(e1 - e2).norm().item() for e2 in embeddings] for e1 in embeddings]
print(names)
print(pd.DataFrame(dists, columns=names, index=names)) 
torch.save(embeddings,'database.pt')  # 当然也可以保存在一个文件
torch.save(names,'names.pt')

有了上述的数据库，就能通过距离识别人脸。这个库的mtcnn模型，只有一个forward和detect方法，其中forward的输出是人脸的3*image_size*image_size的张量，如果参数keep_all = True，则会增加一个维度，返回所有检测到的人脸张量，即N*3*image_size*image_size的张量，另外参数return_prob决定是否返回概率值，即检测到人脸的概率，典型用法如下。而detect方法则返回了检测的人脸框的位置坐标和概率。

因此，forward的输出是NCHW的张量，用于后面的模型输入计算特征向量，而detect的输出是人脸的位置坐标，如果两个结果都需要的话，就得根据例子进行两次运算，相当于进行了两次推理运算。

mtcnn = MTCNN()
face_tensor, prob = mtcnn(img, save_path='face.png', return_prob=True)  # 返回的是检测到的人脸数据tensor 形状是N，3*160*160 # 尺寸不一定是160，之前的参数设置
boxes, prob = mtcnn.detect(img)  # 直接返回人脸的位置坐标和概率

对新的照片进行人脸识别

# mtcnn网络负责检测人脸 
mtcnn = MTCNN(keep_all=True, device=device)
resnet = InceptionResnetV1(pretrained='vggface2').eval().to('cuda')

names = torch.load("./database/names.pt")
embeddings = torch.load("./database/database.pt").to('cuda')
def detect_frame(img):
    fontStyle = ImageFont.truetype("LiberationSans-Regular.ttf", 25,encoding="utf-8")
    faces = mtcnn(img)  # 直接infer所有的faces
    #但是这里相当于两次infer，会浪费时间
    boxes, _ = mtcnn.detect(img)  # 检测出人脸框 返回的是位置 
    frame_draw = img.copy()
    draw = ImageDraw.Draw(frame_draw)
    print("检测人脸数目：",len(boxes))
    for i,box in enumerate(boxes):
        draw.rectangle(box.tolist(), outline=(255, 0, 0))  # 绘制框
        face_embedding = resnet(faces[i].unsqueeze(0).to('cuda'))
        #print(face_embedding.size(),'大小')
        # 计算距离
        probs = [(face_embedding - embeddings[i]).norm().item() for i in range(embeddings.size()[0])] 
        #print(probs)
        # 我们可以认为距离最近的那个就是最有可能的人，但也有可能出问题，数据库中可以存放一个人的多视角多姿态数据，对比的时候可以采用其他方法，如投票机制决定最后的识别人脸
        index = probs.index(min(probs))   # 对应的索引就是判断的人脸
        name = names[index] # 对应的人脸
        draw.text( (int(box[0]),int(box[1])), str(name), fill=(255,0,0),font=fontStyle)
    return frame_draw

上述步骤还有很多地方可以更改，其中如果要想一步完成face_tensor的计算送入分类网络，与同时返回检测到的face_box，可以通过修改源码完成。

通过观察，可以发现在forward的源码第一句实际上就是调用了detect方法，但是最后没有返回检测到的batch_boxes，而是通过它获取图像数据，因此，可以直接在返回值中增加一个目标框即可（但是如果后续要对模型进行调整训练，又会有其他部分的改动，比如训练时的损失计算这些，infer的结果应该是返回值的部分，毕竟原先只有一个返回值）。

def forward(self, img, save_path=None, return_prob=False):
      # 这里是mtcnn的forward源码。可以看到，实际上也是调用了detect函数，返回了框坐标和概率，后面对框进行处理，从而返回实际的人脸数据
    with torch.no_grad():
                batch_boxes, batch_probs = self.detect(img)

另外作者提供了InceptionResnetV1的微调demo，按照数据集的格式组织文件后，基于预训练模型在自己的数据集上调整，可以效果更好。