利用facenet-pytorch实现人脸检测,并提取人脸特征求余弦相似度,代码如下
from facenet_pytorch import MTCNN, InceptionResnetV1
from sklearn.metrics.pairwise import cosine_similarity
from PIL import Image
import numpy as np
# If required, create a face detection pipeline using MTCNN:
#mtcnn = MTCNN(image_size=<image_size>, margin=<margin>)
mtcnn = MTCNN()
# Create an inception resnet (in eval mode):
resnet = InceptionResnetV1(pretrained='vggface2').eval()
img = Image.open("./ldh1.jpeg")
# Get cropped and prewhitened image tensor
#img_cropped = mtcnn(img, save_path=<optional save path>)
img_cropped = mtcnn(img, "./mtcnnresult1.jpg")
# Calculate embedding (unsqueeze to add batch dimension)
img_embedding = resnet(img_cropped.unsqueeze(0))
features1 = img_embedding.detach().numpy()
print("length os features1::", len(features1[0]))
print("features1::", features1[0])
img = Image.open("./ldh2.jpeg")
# Get cropped and prewhitened image tensor
#img_cropped = mtcnn(img, save_path=<optional save path>)
img_cropped = mtcnn(img, "./mtcnnresult2.jpg")
# Calculate embedding (unsqueeze to add batch dimension)
img_embedding = resnet(img_cropped.unsqueeze(0))
features2 = img_embedding.detach().numpy()
print("length os features2::", len(features2[0]))
print("features2::", features2[0])
dist12 = cosine_similarity(features1, features2)
print("dist12 is:", dist12)
# Or, if using for VGGFace2 classification
#resnet.classify = True
#img_probs = resnet(img_cropped.unsqueeze(0))
参考文献:
https://github.com/timesler/facenet-pytorch
|