FaceNet class(见 bubbliiiing-github)
class Facenet(object):
_defaults = {
#--------------------------------------------------------------------------#
# 使用自己训练好的模型进行预测要修改model_path,指向logs文件夹下的权值文件
# 训练好后logs文件夹下存在多个权值文件,选择验证集损失较低的即可。
# 验证集损失较低不代表准确度较高,仅代表该权值在验证集上泛化性能较好。
#--------------------------------------------------------------------------#
"model_path" : "model_data/facenet_inception_resnetv1.pth",
#--------------------------------------------------------------------------#
# 输入图片的大小。
#--------------------------------------------------------------------------#
"input_shape" : [160, 160, 3],
#--------------------------------------------------------------------------#
# 所使用到的主干特征提取网络
#--------------------------------------------------------------------------#
"backbone" : "inception_resnetv1",
#-------------------------------------------#
# 是否进行不失真的resize
#-------------------------------------------#
"letterbox_image" : True,
#-------------------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------------------#
"cuda" : True,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Facenet
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
for name, value in kwargs.items():
setattr(self, name, value)
self.generate()
def generate(self):
#---------------------------------------------------#
# 载入模型与权值
#---------------------------------------------------#
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.net = facenet(backbone=self.backbone, mode="predict").eval()
self.net.load_state_dict(torch.load(self.model_path, map_location=device), strict=False)
print('{} model loaded.'.format(self.model_path))
if self.cuda:
self.net = torch.nn.DataParallel(self.net)
cudnn.benchmark = True
self.net = self.net.cuda()
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image_1, image_2,mode):
#---------------------------------------------------#
# 图片预处理,归一化
#---------------------------------------------------#
with torch.no_grad():
image_1 = resize_image(image_1, [self.input_shape[1], self.input_shape[0]], letterbox_image=self.letterbox_image)
image_2 = resize_image(image_2, [self.input_shape[1], self.input_shape[0]], letterbox_image=self.letterbox_image)
photo_1 = torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image_1, np.float32)), (2, 0, 1)), 0))
photo_2 = torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image_2, np.float32)), (2, 0, 1)), 0))
if self.cuda:
photo_1 = photo_1.cuda()
photo_2 = photo_2.cuda()
#---------------------------------------------------#
# 图片传入网络进行预测
#---------------------------------------------------#
output1 = self.net(photo_1).cpu().numpy()
output2 = self.net(photo_2).cpu().numpy()
#---------------------------------------------------#
# 计算二者之间的距离
#---------------------------------------------------#
l1 = np.linalg.norm(output1 - output2, axis=1)
if mode == True:
plt.subplot(1, 2, 1)
plt.imshow(np.array(image_1))
plt.subplot(1, 2, 2)
plt.imshow(np.array(image_2))
plt.text(-12, -12, 'Distance:%.3f' % l1, ha='center', va= 'bottom',fontsize=11)
plt.show()
return l1
添加 Face match 功能,一对多的匹配。
import os
from PIL import Image
from facenet import Facenet
def Face_match(face_path, face_database, mode):
"""
:param face_path: 放需要检测的人脸
:param face_database: 匹配的人脸库
:param mode: True/False 是否需要展示匹配的图片(同时显示distance)
:return:
"""
model = Facenet ()
face_unkonw_path = face_path
face_known_path = face_database
face_unknown = []
face_database = []
for file in os.listdir (face_unkonw_path):
face_unknown.append (file)
for file in os.listdir (face_known_path):
face_database.append (file)
print ('-' * 60)
print ("Unknowm faces: {} ".format (face_unknown))
print ("Face database: {} ".format (face_database))
print ('-' * 60)
for unkonw_face in face_unknown:
face_distance = []
face_path_1 = os.path.join (face_unkonw_path, unkonw_face)
face_img_1 = Image.open (face_path_1)
for face in face_database:
face_path_2 = os.path.join (face_known_path, face)
face_img_2 = Image.open (face_path_2)
probability = model.detect_image (face_img_1, face_img_2, mode)
face_distance.append (probability[0])
index = face_distance.index (min (face_distance))
print ("The matched face for {} is: {} ".format (unkonw_face, face_database[index]))
if __name__ == "__main__":
face_unkonw_path = r'FaceImages'
face_known_path = r'FaceDatabase'
mode = False
Face_match (face_unkonw_path, face_known_path, mode)