378 lines
14 KiB
Python
378 lines
14 KiB
Python
import os
|
|
import time
|
|
import re
|
|
import torch
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from anti import anti_spoofing, load_anti_model
|
|
from backbones import iresnet50, iresnet18, iresnet100
|
|
from retinaface_detect import load_retinaface_model, detect_one, detect_video, set_retinaface_conf
|
|
from torch2trt import torch2trt, TRTModule
|
|
|
|
threshold = 0.7
|
|
|
|
|
|
# 读取112x112的本地图片并变换通道位置归一化
|
|
def load_image(img_path):
|
|
img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
|
|
img = img.transpose((2, 0, 1))
|
|
img = img[np.newaxis, :, :, :]
|
|
img = np.array(img, dtype=np.float32)
|
|
img -= 127.5
|
|
img /= 127.5
|
|
return img
|
|
|
|
|
|
# 计算两个特征向量的欧式距离
|
|
def findEuclideanDistance(source_representation, test_representation):
|
|
euclidean_distance = source_representation - test_representation
|
|
euclidean_distance = np.sum(np.multiply(euclidean_distance, euclidean_distance))
|
|
euclidean_distance = np.sqrt(euclidean_distance)
|
|
return euclidean_distance
|
|
|
|
|
|
# 计算两个特征向量的余弦距离
|
|
def findCosineDistance(source_representation, test_representation):
|
|
a = np.matmul(np.transpose(source_representation), test_representation)
|
|
b = np.sum(np.multiply(source_representation, source_representation))
|
|
c = np.sum(np.multiply(test_representation, test_representation))
|
|
return 1 - (a / (np.sqrt(b) * np.sqrt(c)))
|
|
|
|
|
|
# 归一化欧氏距离
|
|
def l2_normalize(x):
|
|
return x / np.sqrt(np.sum(np.multiply(x, x)))
|
|
|
|
|
|
# 归一化余弦距离
|
|
def cosin_metric(x1, x2):
|
|
return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2))
|
|
|
|
|
|
# 加载保存的姓名、人脸特征向量的人脸库
|
|
def load_npy(path):
|
|
data = np.load(path, allow_pickle=True)
|
|
data = data.item()
|
|
return data
|
|
|
|
|
|
# 批量化生成人脸特征向量并保存到人脸库
|
|
def create_database_batch(path, model, database_path):
|
|
name_list = os.listdir(path)
|
|
k_v = {}
|
|
if os.path.exists(database_path):
|
|
k_v = np.load(database_path, allow_pickle=True)
|
|
k_v = k_v.item()
|
|
batch = 256
|
|
order_name = []
|
|
order_path = []
|
|
emb_list = []
|
|
for name in name_list[:]:
|
|
img_path = os.path.join(path, name)
|
|
# for img_name in img_path[:1]:
|
|
order_name.append(name[:-4])
|
|
order_path.append(img_path)
|
|
order_img = np.zeros((len(order_path), 3, 112, 112), dtype=np.float32)
|
|
for index, img_path in enumerate(order_path):
|
|
order_img[index] = load_image(img_path)
|
|
print(order_img.shape)
|
|
order_img = torch.from_numpy(order_img)
|
|
order_img = order_img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
now = 0
|
|
number = len(order_img)
|
|
with torch.no_grad():
|
|
while now < number:
|
|
if now + batch < number:
|
|
emb = model(order_img[now:now + batch])
|
|
else:
|
|
emb = model(order_img[now:])
|
|
now = now + batch
|
|
emb = emb.cpu().numpy()
|
|
for em in emb:
|
|
emb_list.append(em)
|
|
print("batch" + str(now))
|
|
|
|
for i, emb in enumerate(emb_list):
|
|
k_v[order_name[i]] = l2_normalize(emb)
|
|
np.save(database_path, k_v)
|
|
|
|
def create_database_from_img(order_name, order_img, model, database_path, cpu_or_cuda):
|
|
k_v = {}
|
|
if os.path.exists(database_path):
|
|
k_v = np.load(database_path, allow_pickle=True)
|
|
k_v = k_v.item()
|
|
batch = 256
|
|
emb_list = []
|
|
|
|
print(order_img.shape)
|
|
order_img = torch.from_numpy(order_img)
|
|
order_img = order_img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
now = 0
|
|
number = len(order_img)
|
|
with torch.no_grad():
|
|
while now < number:
|
|
if now + batch < number:
|
|
emb = model(order_img[now:now + batch])
|
|
else:
|
|
emb = model(order_img[now:])
|
|
now = now + batch
|
|
emb = emb.cpu().numpy()
|
|
for em in emb:
|
|
emb_list.append(em)
|
|
print("batch" + str(now))
|
|
for i, emb in enumerate(emb_list):
|
|
k_v[order_name[i]] = l2_normalize(emb)
|
|
np.save(database_path, k_v)
|
|
|
|
# 向人脸库中新增一个人的姓名和人脸特征向量，若人脸库不存在则创建
|
|
def add_one_to_database(img, model, name, database_path, cpu_or_cuda):
|
|
img = torch.from_numpy(img)
|
|
img = img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
with torch.no_grad():
|
|
pred = model(img)
|
|
pred = pred.cpu().numpy()
|
|
k_v = {}
|
|
if os.path.exists(database_path):
|
|
k_v = np.load(database_path, allow_pickle=True)
|
|
k_v = k_v.item()
|
|
k_v[name] = l2_normalize(pred)
|
|
np.save(database_path, k_v)
|
|
|
|
|
|
# 计算此特征向量与人脸库中的哪个人脸特征向量距离最近
|
|
def findmindistance(pred, threshold, k_v):
|
|
distance = 10
|
|
most_like = ""
|
|
for name in k_v.keys():
|
|
tmp = findEuclideanDistance(k_v[name], pred)
|
|
if distance > tmp:
|
|
distance = tmp
|
|
most_like = name
|
|
if distance < threshold:
|
|
return most_like, distance
|
|
else:
|
|
return -1, distance
|
|
|
|
|
|
def faiss_find_face(pred, index, database_name_list):
|
|
name_list = []
|
|
start_time = time.time()
|
|
D, I = index.search(pred, 1)
|
|
end_time = time.time()
|
|
# print("faiss cost %fs" % (end_time - start_time))
|
|
# print(D, I)
|
|
if len(pred) == 1:
|
|
if D[0][0] < threshold:
|
|
# print(database_name_list[I[0][0]])
|
|
return database_name_list[I[0][0]], D[0][0]
|
|
else:
|
|
return "unknown", D[0][0]
|
|
else:
|
|
for i,index in enumerate(I):
|
|
if D[i][0] < threshold:
|
|
#print(database_name_list[I[0][0]])
|
|
name_list.append(database_name_list[index[0]]+str(D[i][0]))
|
|
else:
|
|
name_list.append("unknown"+str(D[i][0]))
|
|
return name_list
|
|
|
|
|
|
# 从人脸库中找到单个人脸
|
|
def findOne(img, model, index, database_name_list, cpu_or_cuda):
|
|
img = torch.from_numpy(img)
|
|
img = img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
with torch.no_grad():
|
|
start_time = time.time()
|
|
pred = model(img)
|
|
end_time = time.time()
|
|
print("predOne time: " + str(end_time - start_time))
|
|
pred = pred.cpu().numpy()
|
|
# start_time = time.time()
|
|
# name, distance = findmindistance(l2_normalize(pred), threshold=threshold, k_v=k_v)
|
|
# end_time = time.time()
|
|
# print("baoli time: " + str(end_time - start_time))
|
|
name, distance = faiss_find_face(l2_normalize(pred), index, database_name_list)
|
|
print(pred.shape)
|
|
if name != -1:
|
|
mo = r'[\u4e00-\u9fa5_a-zA-Z0-9]*'
|
|
name = re.match(mo, name)
|
|
return name.group(0), distance
|
|
else:
|
|
return "unknown", distance
|
|
|
|
|
|
# 从人脸库中找到传入的人脸列表中的所有人脸
|
|
def findAll(imglist, model, index ,database_name_list, cpu_or_cuda):
|
|
imglist = torch.from_numpy(imglist)
|
|
imglist = imglist.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
with torch.no_grad():
|
|
name_list =[]
|
|
start_time = time.time()
|
|
pred = model(imglist)
|
|
end_time = time.time()
|
|
print("predOne time: " + str(end_time - start_time))
|
|
pred = pred.cpu().numpy()
|
|
start_time = time.time()
|
|
#name_list = faiss_find_face(l2_normalize(pred), index, database_name_list)
|
|
for pr in pred:
|
|
pr = np.expand_dims(l2_normalize(pr), 0)
|
|
# #print(pr.shape)
|
|
name, distance = faiss_find_face(l2_normalize(pr), index, database_name_list)
|
|
#name_list.append(name+" "+str(distance))
|
|
name_list.append(name)
|
|
# for pr in pred:
|
|
# name, distance = findmindistance(l2_normalize(pr), threshold=threshold, k_v=k_v)
|
|
# if name != -1:
|
|
# mo = r'[\u4e00-\u9fa5_a-zA-Z]*'
|
|
# name = re.match(mo, name)
|
|
# name_list.append(name.group(0) + str(distance))
|
|
# else:
|
|
# name_list.append("unknown" + str(distance))
|
|
end_time = time.time()
|
|
print("searchALL time: " + str(end_time - start_time))
|
|
return name_list
|
|
|
|
|
|
# 提取为512维特征向量
|
|
def embedding(order_img, model, cpu_or_cuda):
|
|
number = len(order_img)
|
|
order_img = torch.from_numpy(order_img)
|
|
order_img = order_img.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
batch = 64
|
|
emb_list = []
|
|
now = 0
|
|
with torch.no_grad():
|
|
while now < number:
|
|
if now + batch < number:
|
|
emb = model(order_img[now:now + batch])
|
|
else:
|
|
emb = model(order_img[now:])
|
|
now = now + batch
|
|
emb = emb.cpu().numpy()
|
|
for em in emb:
|
|
emb_list.append(l2_normalize(em))
|
|
# print("batch" + str(now))
|
|
emb_list = np.array(emb_list)
|
|
return emb_list
|
|
|
|
|
|
# 处理聚类人脸文件夹，返回特征向量列表，文件名列表
|
|
def get_claster_tmp_file_embedding(file_path, retinaface_model, retinaface_args, arcface_model, cpu_or_cuda):
|
|
img_name = os.listdir(file_path)
|
|
img_list = []
|
|
for name in img_name:
|
|
all_face, box_and_point = detect_one(os.path.join(file_path, name), retinaface_model, retinaface_args)
|
|
img_list.append(all_face[0])
|
|
img_list = np.array(img_list)
|
|
# print(img_list.shape)
|
|
emb_list = embedding(img_list, arcface_model, cpu_or_cuda)
|
|
return emb_list, img_name
|
|
|
|
|
|
# 同一个人聚为一类
|
|
def cluster(emb_list, name_list):
|
|
all_claster = []
|
|
cla = []
|
|
in_claster_name = []
|
|
img_number = len(emb_list)
|
|
for index, emb in enumerate(emb_list):
|
|
if name_list[index] in in_claster_name:
|
|
continue
|
|
for j in range(img_number - index - 1):
|
|
if findEuclideanDistance(emb, emb_list[index + 1 + j]) < threshold:
|
|
if name_list[index + 1 + j] not in in_claster_name:
|
|
cla.append(name_list[index + 1 + j])
|
|
in_claster_name.append(name_list[index + 1 + j])
|
|
cla.append(name_list[index])
|
|
in_claster_name.append(name_list[index])
|
|
all_claster.append(cla)
|
|
cla = []
|
|
return all_claster
|
|
|
|
|
|
# 加载人脸识别模型
|
|
def load_arcface_model(model_path, cpu_or_cuda):
|
|
if cpu_or_cuda == "trt":
|
|
model = TRTModule()
|
|
model.load_state_dict(torch.load('./model/arcface_trt.pth'))
|
|
elif cpu_or_cuda == "trt_new":
|
|
model = iresnet100()
|
|
model.load_state_dict(torch.load(model_path, map_location="cuda"))
|
|
model = model.eval()
|
|
model.to(torch.device("cuda"))
|
|
x = torch.ones((1, 3, 112, 112)).to(torch.device("cuda"))
|
|
model = torch2trt(model, [x], max_batch_size=4)
|
|
torch.save(model.state_dict(), './model/arcface_trt.pth')
|
|
else:
|
|
model = iresnet100()
|
|
model.load_state_dict(torch.load(model_path, map_location=cpu_or_cuda))
|
|
model = model.eval()
|
|
model.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
return model
|
|
|
|
|
|
# 对比两张人脸是否相同
|
|
def face_verification(img1, img2, model, cpu_or_cuda):
|
|
img_list = np.concatenate((img1, img2), axis=0)
|
|
img_list = torch.from_numpy(img_list)
|
|
img_list = img_list.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
|
with torch.no_grad():
|
|
pred = model(img_list)
|
|
pred = pred.cpu().numpy()
|
|
distance = findEuclideanDistance(l2_normalize(pred[0]), l2_normalize(pred[1]))
|
|
# print("EuclideanDistance is :" + str(distance))
|
|
if distance < threshold:
|
|
return 'same ',distance
|
|
else:
|
|
return 'different ', distance
|
|
|
|
|
|
if __name__ == '__main__':
|
|
cpu_or_cuda = "cuda" if torch.cuda.is_available() else "cpu"
|
|
arcface_model = load_arcface_model("./model/backbone100.pth", cpu_or_cuda=cpu_or_cuda)
|
|
# retinaface_args = set_retinaface_conf(cpu_or_cuda=cpu_or_cuda)
|
|
# retinaface_model = load_retinaface_model(retinaface_args)
|
|
#
|
|
# anti_spoofing_model_path = "model/anti_spoof_models"
|
|
# anti_model = load_anti_model(anti_spoofing_model_path, 0)
|
|
#
|
|
# k_v = load_npy("./Database/student.npy")
|
|
# 对比两张人脸
|
|
# img1, box_and_point = detect_one("D:\Download\lfw\lfw\Aaron_Peirsol\Aaron_Peirsol_0001.jpg", retinaface_model, retinaface_args)
|
|
# img2, box_and_point = detect_one("D:\Download\lfw\lfw\Aaron_Peirsol\Aaron_Peirsol_0002.jpg", retinaface_model, retinaface_args)
|
|
# print(face_verification(img1, img2, arcface_model))
|
|
|
|
# img3 = load_image(r"D:\Download\out\alig_students\student.jpg")
|
|
# img3 = torch.from_numpy(img3)
|
|
# 单张人脸活体检测
|
|
# img3, b_p = detect_one(r"C:\Users\ASUS\Desktop\face\IMG_20210525_113950.jpg", retinaface_model, retinaface_args)
|
|
# b = b_p[0]
|
|
# w = b[2] - b[0]
|
|
# h = b[3] - b[1]
|
|
# b[2] = w
|
|
# b[3] = h
|
|
# label, value = anti_spoofing("./img/recognition/000_0.bmp", "model/anti_spoof_models", 0, np.array(b[:4], int), anti_model)
|
|
# print(label,value)
|
|
# name = findOne(img3, arcface_model, k_v, cpu_or_cuda)
|
|
# print(name)
|
|
|
|
# 人脸聚类
|
|
# emb_list, name_list = get_claster_tmp_file_embedding("./img/cluster_tmp_file/face", retinaface_model,
|
|
# retinaface_args, arcface_model, cpu_or_cuda)
|
|
# print(cluster(emb_list, name_list))
|
|
|
|
# img3, box_and_point = detect_one("D:\Download\out\students\student.jpg", retinaface_model, retinaface_args)
|
|
# print(embedding(img3,arcface_model).shape)
|
|
|
|
# 人脸库中增加一张人脸
|
|
# add_one_to_database(img1,arcface_model,"Aaron_Peirsol","./Database/student.npy")
|
|
# name = findOne(img1, arcface_model, k_v)
|
|
# print(name)
|
|
|
|
# 人脸库中批量增加人脸
|
|
create_database_batch(r"D:\Download\out\alig_students_all", arcface_model, "./Database/sfz.npy")
|
|
|
|
# 识别视频中的人脸
|
|
# detect_video("software.mp4","out.avi",retinaface_model,arcface_model,k_v,retinaface_args)
|