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,按照数据集的格式组织文件后,基于预训练模型在自己的数据集上调整,可以效果更好。