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，按照數據集的格式組織文件后，基於預訓練模型在自己的數據集上調整，可以效果更好。