484 lines
18 KiB
Python
484 lines
18 KiB
Python
from __future__ import print_function
|
||
import re
|
||
import time
|
||
import cv2
|
||
import torch
|
||
import torch.backends.cudnn as cudnn
|
||
import numpy as np
|
||
from skimage import transform as trans
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
from data import cfg_mnet, cfg_re50
|
||
from layers.functions.prior_box import PriorBox
|
||
from utils.nms.py_cpu_nms import py_cpu_nms
|
||
from models.retinaface import RetinaFace
|
||
from utils.box_utils import decode, decode_landm
|
||
|
||
threshold = 1.05
|
||
ppi = 1280
|
||
step = 3
|
||
|
||
class ConfRetinaface(object):
|
||
def __init__(self, trained_model, network, cpu, confidence_threshold, top_k, nms_threshold, keep_top_k, vis_thres):
|
||
self.trained_model = trained_model
|
||
self.network = network
|
||
self.cpu = cpu
|
||
self.confidence_threshold = confidence_threshold
|
||
self.top_k = top_k
|
||
self.nms_threshold = nms_threshold
|
||
self.keep_top_k = keep_top_k
|
||
self.vis_thres = vis_thres
|
||
|
||
|
||
def set_retinaface_conf(cpu_or_cuda):
|
||
args = ConfRetinaface(trained_model='./weights/mobilenet0.25_Final.pth',
|
||
network='mobile0.25',
|
||
cpu=True if cpu_or_cuda == 'cpu' else False,
|
||
confidence_threshold=0.02,
|
||
top_k=5000,
|
||
nms_threshold=0.4,
|
||
keep_top_k=750,
|
||
vis_thres=0.6)
|
||
return args
|
||
|
||
|
||
def check_keys(model, pretrained_state_dict):
|
||
ckpt_keys = set(pretrained_state_dict.keys())
|
||
model_keys = set(model.state_dict().keys())
|
||
used_pretrained_keys = model_keys & ckpt_keys
|
||
unused_pretrained_keys = ckpt_keys - model_keys
|
||
missing_keys = model_keys - ckpt_keys
|
||
print('Missing keys:{}'.format(len(missing_keys)))
|
||
print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))
|
||
print('Used keys:{}'.format(len(used_pretrained_keys)))
|
||
assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
|
||
return True
|
||
|
||
|
||
def remove_prefix(state_dict, prefix):
|
||
''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''
|
||
print('remove prefix \'{}\''.format(prefix))
|
||
f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
|
||
return {f(key): value for key, value in state_dict.items()}
|
||
|
||
|
||
def load_model(model, pretrained_path, load_to_cpu):
|
||
print('Loading pretrained model from {}'.format(pretrained_path))
|
||
if load_to_cpu:
|
||
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
|
||
else:
|
||
device = torch.cuda.current_device()
|
||
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
|
||
if "state_dict" in pretrained_dict.keys():
|
||
pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')
|
||
else:
|
||
pretrained_dict = remove_prefix(pretrained_dict, 'module.')
|
||
check_keys(model, pretrained_dict)
|
||
model.load_state_dict(pretrained_dict, strict=False)
|
||
return model
|
||
|
||
|
||
# 加载retinaface模型
|
||
def load_retinaface_model(args):
|
||
torch.set_grad_enabled(False)
|
||
cfg = None
|
||
if args.network == "mobile0.25":
|
||
cfg = cfg_mnet
|
||
elif args.network == "resnet50":
|
||
cfg = cfg_re50
|
||
# net and model
|
||
net = RetinaFace(cfg=cfg, phase='test')
|
||
net = load_model(net, args.trained_model, args.cpu)
|
||
net.eval()
|
||
cudnn.benchmark = True
|
||
device = torch.device("cpu" if args.cpu else "cuda")
|
||
net = net.to(device)
|
||
print('Finished loading model!')
|
||
return net
|
||
|
||
|
||
# 计算两个特征向量的欧式距离
|
||
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 l2_normalize(x):
|
||
return x / np.sqrt(np.sum(np.multiply(x, x)))
|
||
|
||
|
||
# 计算此特征向量与人脸库中的哪个人脸特征向量距离最近
|
||
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
|
||
else:
|
||
return "unknown"
|
||
|
||
#
|
||
def faiss_find_face(pred,index ,database_name_list):
|
||
#print(len(database_name_list))
|
||
start_time = time.time()
|
||
D, I = index.search(pred, 1)
|
||
name_list = []
|
||
end_time = time.time()
|
||
print("faiss cost %fs" % (end_time - start_time))
|
||
print(D, I)
|
||
# if D[0][0] < threshold:
|
||
# print(database_name_list[I[0][0]])
|
||
# return database_name_list[I[0][0]]
|
||
# else:
|
||
# return "unknown"
|
||
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]])
|
||
else:
|
||
name_list.append("unknown")
|
||
return name_list
|
||
|
||
# 从人脸库中找到传入的人脸列表中的所有人脸
|
||
def findAll(imglist, model, index ,database_name_list, k_v, cpu_or_cuda):
|
||
start_time = time.time()
|
||
imglist = torch.from_numpy(imglist)
|
||
imglist = imglist.to(torch.device("cuda" if cpu_or_cuda == "cuda" else "cpu"))
|
||
with torch.no_grad():
|
||
name_list = []
|
||
pred = model(imglist)
|
||
pred = pred.cpu().numpy()
|
||
print("predOne time: " + str(time.time() - start_time))
|
||
#print(pred.shape)
|
||
start_time = time.time()
|
||
#name_list = faiss_find_face(l2_normalize(pred), index, database_name_list)
|
||
for pr in pred:
|
||
name = findmindistance(l2_normalize(pr), threshold=threshold, k_v=k_v)
|
||
print(name)
|
||
# print(l2_normalize(pr).shape)
|
||
#pr = np.expand_dims(l2_normalize(pr), 0)
|
||
#print(pr.shape)
|
||
#name = faiss_find_face(pr,index ,database_name_list)
|
||
if name != "unknown":
|
||
mo = r'[\u4e00-\u9fa5_a-zA-Z]*'
|
||
name = re.match(mo, name)
|
||
name_list.append(name.group(0))
|
||
else:
|
||
name_list.append("unknown")
|
||
#name_list.append(name)
|
||
print("findOne time: " + str(time.time() - start_time))
|
||
return name_list
|
||
|
||
|
||
# 检测单张人脸，返回1x3x112x112的数组
|
||
def detect_one(path, net, args):
|
||
cfg = None
|
||
if args.network == "mobile0.25":
|
||
cfg = cfg_mnet
|
||
elif args.network == "resnet50":
|
||
cfg = cfg_re50
|
||
|
||
device = torch.device("cpu" if args.cpu else "cuda")
|
||
resize = 1
|
||
|
||
# testing begin
|
||
frame = cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)
|
||
h, w = frame.shape[:2]
|
||
factor = h / w
|
||
if (w > 1000):
|
||
frame = cv2.resize(frame, (600, int(600 * factor)))
|
||
h, w = frame.shape[:2]
|
||
|
||
tic = time.time()
|
||
img = np.float32(frame)
|
||
im_height, im_width, _ = img.shape
|
||
scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
|
||
img -= (104, 117, 123)
|
||
img = img.transpose(2, 0, 1)
|
||
img = torch.from_numpy(img).unsqueeze(0)
|
||
img = img.to(device)
|
||
scale = scale.to(device)
|
||
|
||
loc, conf, landms = net(img) # forward pass
|
||
#print(loc.shape,landms.shape,conf.shape)
|
||
priorbox = PriorBox(cfg, image_size=(im_height, im_width))
|
||
priors = priorbox.forward()
|
||
priors = priors.to(device)
|
||
prior_data = priors.data
|
||
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
|
||
boxes = boxes * scale / resize
|
||
boxes = boxes.cpu().numpy()
|
||
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
|
||
landms = decode_landm(landms.data.squeeze(0), prior_data, cfg['variance'])
|
||
scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||
img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||
img.shape[3], img.shape[2]])
|
||
scale1 = scale1.to(device)
|
||
landms = landms * scale1 / resize
|
||
landms = landms.cpu().numpy()
|
||
|
||
# ignore low scores
|
||
inds = np.where(scores > args.confidence_threshold)[0]
|
||
boxes = boxes[inds]
|
||
landms = landms[inds]
|
||
scores = scores[inds]
|
||
|
||
# keep top-K before NMS
|
||
order = scores.argsort()[::-1][:args.top_k]
|
||
boxes = boxes[order]
|
||
landms = landms[order]
|
||
scores = scores[order]
|
||
|
||
# do NMS
|
||
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
|
||
keep = py_cpu_nms(dets, args.nms_threshold)
|
||
# keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
|
||
dets = dets[keep, :]
|
||
landms = landms[keep]
|
||
|
||
# keep top-K faster NMS
|
||
dets = dets[:args.keep_top_k, :]
|
||
landms = landms[:args.keep_top_k, :]
|
||
|
||
dets = np.concatenate((dets, landms), axis=1)
|
||
face_list = []
|
||
box_and_point = []
|
||
# print(dets[:4])
|
||
# print('net forward time: {:.4f}'.format(time.time() - tic))
|
||
print(len(dets))
|
||
for i, det in enumerate(dets):
|
||
|
||
if det[4] < args.vis_thres:
|
||
continue
|
||
box_and_point.append(det)
|
||
dst = np.reshape(landms[i], (5, 2))
|
||
# print(dst.shape)
|
||
src1 = np.array([
|
||
[38.3814, 51.6963],
|
||
[73.6186, 51.5014],
|
||
[56.1120, 71.7366],
|
||
[41.6361, 92.3655],
|
||
[70.8167, 92.2041]], dtype=np.float32)
|
||
# print(src1.shape)
|
||
tform = trans.SimilarityTransform()
|
||
tform.estimate(dst, src1)
|
||
M = tform.params[0:2, :]
|
||
frame2 = cv2.warpAffine(frame, M, (w, h), borderValue=0.0)
|
||
img112 = frame2[0:112, 0:112, :]
|
||
# cv2.imshow('out', img112)
|
||
# cv2.waitKey(0)
|
||
face_list.append(img112)
|
||
if len(face_list) > 0:
|
||
face_list = np.array(face_list)
|
||
face_list = face_list.transpose((0, 3, 1, 2))
|
||
face_list = np.array(face_list, dtype=np.float32)
|
||
face_list -= 127.5
|
||
face_list /= 127.5
|
||
box_and_point = np.array(box_and_point)
|
||
# face_list = torch.from_numpy(face_list)
|
||
# cv2.imshow('out', img112)
|
||
# cv2.waitKey(0)
|
||
return face_list, box_and_point
|
||
|
||
|
||
# 检测视频中的人脸并人脸识别
|
||
def detect_video(video_path, output_path, net, arcface_model, k_v, args):
|
||
tic_total = time.time()
|
||
cfg = None
|
||
if args.network == "mobile0.25":
|
||
cfg = cfg_mnet
|
||
elif args.network == "resnet50":
|
||
cfg = cfg_re50
|
||
device = torch.device("cpu" if args.cpu else "cuda")
|
||
resize = 1
|
||
|
||
# testing begin
|
||
cap = cv2.VideoCapture(video_path)
|
||
ret, frame = cap.read()
|
||
h, w = frame.shape[:2]
|
||
factor = 0
|
||
if (w > ppi):
|
||
factor = h / w
|
||
frame = cv2.resize(frame, (ppi, int(ppi * factor)))
|
||
h, w = frame.shape[:2]
|
||
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
size = (w, h)
|
||
# size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
|
||
# int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
|
||
# out = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps, size)
|
||
out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'XVID'), fps, size)
|
||
number = step
|
||
dets = []
|
||
name_list = []
|
||
font = ImageFont.truetype("font.ttf", 22)
|
||
priorbox = PriorBox(cfg, image_size=(h, w))
|
||
priors = priorbox.forward()
|
||
priors = priors.to(device)
|
||
prior_data = priors.data
|
||
|
||
scale = torch.Tensor([w, h, w, h])
|
||
scale = scale.to(device)
|
||
scale1 = torch.Tensor([w, h, w, h,
|
||
w, h, w, h,
|
||
w, h])
|
||
scale1 = scale1.to(device)
|
||
|
||
src1 = np.array([
|
||
[38.3814, 51.6963],
|
||
[73.6186, 51.5014],
|
||
[56.1120, 71.7366],
|
||
[41.6361, 92.3655],
|
||
[70.8167, 92.2041]], dtype=np.float32)
|
||
# print(src1.shape)
|
||
tform = trans.SimilarityTransform()
|
||
|
||
while ret:
|
||
tic_all = time.time()
|
||
if number == step:
|
||
tic = time.time()
|
||
img = np.float32(frame)
|
||
img -= (104, 117, 123)
|
||
img = img.transpose(2, 0, 1)
|
||
img = torch.from_numpy(img).unsqueeze(0)
|
||
img = img.to(device)
|
||
|
||
loc, conf, landms = net(img) # forward pass
|
||
|
||
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
|
||
boxes = boxes * scale / resize
|
||
boxes = boxes.cpu().numpy()
|
||
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
|
||
landms = decode_landm(landms.data.squeeze(0), prior_data, cfg['variance'])
|
||
|
||
landms = landms * scale1 / resize
|
||
landms = landms.cpu().numpy()
|
||
|
||
# ignore low scores
|
||
inds = np.where(scores > args.confidence_threshold)[0]
|
||
boxes = boxes[inds]
|
||
landms = landms[inds]
|
||
scores = scores[inds]
|
||
|
||
# keep top-K before NMS
|
||
order = scores.argsort()[::-1][:args.top_k]
|
||
boxes = boxes[order]
|
||
landms = landms[order]
|
||
scores = scores[order]
|
||
|
||
# do NMS
|
||
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
|
||
keep = py_cpu_nms(dets, args.nms_threshold)
|
||
# keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
|
||
dets = dets[keep, :]
|
||
landms = landms[keep]
|
||
|
||
# keep top-K faster NMS
|
||
dets = dets[:args.keep_top_k, :]
|
||
landms = landms[:args.keep_top_k, :]
|
||
|
||
dets = np.concatenate((dets, landms), axis=1)
|
||
face_list = []
|
||
name_list = []
|
||
# print(dets[:4])
|
||
print('net forward time: {:.4f}'.format(time.time() - tic))
|
||
start_time = time.time()
|
||
for i, det in enumerate(dets[:4]):
|
||
if det[4] < args.vis_thres:
|
||
continue
|
||
boxes, score = det[:4], det[4]
|
||
dst = np.reshape(landms[i], (5, 2))
|
||
# print(dst.shape)
|
||
|
||
tform.estimate(dst, src1)
|
||
M = tform.params[0:2, :]
|
||
frame2 = cv2.warpAffine(frame, M, (w, h), borderValue=0.0)
|
||
img112 = frame2[0:112, 0:112, :]
|
||
face_list.append(img112)
|
||
|
||
if len(face_list) != 0:
|
||
face_list = np.array(face_list)
|
||
face_list = face_list.transpose((0, 3, 1, 2))
|
||
face_list = np.array(face_list, dtype=np.float32)
|
||
face_list -= 127.5
|
||
face_list /= 127.5
|
||
print(face_list.shape)
|
||
# face_list = torch.from_numpy(face_list)
|
||
name_list = findAll(face_list, arcface_model, k_v, "cpu" if args.cpu else "cuda")
|
||
end_time = time.time()
|
||
print("findOneframe time: " + str(end_time - start_time))
|
||
start_time = time.time()
|
||
if (len(dets) != 0):
|
||
for i, det in enumerate(dets[:4]):
|
||
if det[4] < args.vis_thres:
|
||
continue
|
||
boxes, score = det[:4], det[4]
|
||
cv2.rectangle(frame, (int(boxes[0]), int(boxes[1])), (int(boxes[2]), int(boxes[3])), (2, 255, 0), 1)
|
||
|
||
# if (len(dets) != 0):
|
||
# img_PIL = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||
# draw = ImageDraw.Draw(img_PIL)
|
||
#
|
||
# for i, det in enumerate(dets[:4]):
|
||
# if det[4] < args.vis_thres:
|
||
# continue
|
||
# boxes, score = det[:4], det[4]
|
||
# # print(name_list)
|
||
# name = name_list[i]
|
||
# if not isinstance(name, np.unicode):
|
||
# name = name.decode('utf8')
|
||
# draw.text((int(boxes[0]), int(boxes[1])), name, fill=(255, 0, 0), font=font)
|
||
# draw.rectangle((int(boxes[0]), int(boxes[1]), int(boxes[2]), int(boxes[3])), outline="green", width=3)
|
||
# frame = cv2.cvtColor(np.asarray(img_PIL), cv2.COLOR_RGB2BGR)
|
||
#cv2.imshow('out', frame)
|
||
#cv2.waitKey(0)
|
||
out.write(frame)
|
||
end_time = time.time()
|
||
print("drawOneframe time: " + str(end_time - start_time))
|
||
# Press Q on keyboard to stop recording
|
||
# if cv2.waitKey(1) & 0xFF == ord('q'):
|
||
# break
|
||
ret, frame = cap.read()
|
||
number = 0
|
||
if (ret != 0 and factor != 0):
|
||
frame = cv2.resize(frame, (ppi, int(ppi * factor)))
|
||
else:
|
||
number += 1
|
||
if (len(dets) != 0):
|
||
img_PIL = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||
draw = ImageDraw.Draw(img_PIL)
|
||
for i, det in enumerate(dets[:4]):
|
||
if det[4] < args.vis_thres:
|
||
continue
|
||
boxes, score = det[:4], det[4]
|
||
# print(name_list)
|
||
name = name_list[i]
|
||
if not isinstance(name, np.unicode):
|
||
name = name.decode('utf8')
|
||
draw.text((int(boxes[0]), int(boxes[1])), name, fill=(255, 0, 0), font=font)
|
||
draw.rectangle((int(boxes[0]), int(boxes[1]), int(boxes[2]), int(boxes[3])), outline="green",
|
||
width=3)
|
||
frame = cv2.cvtColor(np.asarray(img_PIL), cv2.COLOR_RGB2BGR)
|
||
out.write(frame)
|
||
start_time = time.time()
|
||
ret, frame = cap.read()
|
||
if (ret != 0 and factor != 0):
|
||
frame = cv2.resize(frame, (ppi, int(ppi * factor)))
|
||
print("readframe time: " + str(time.time() - start_time))
|
||
print('all time: {:.4f}'.format(time.time() - tic_all))
|
||
cap.release()
|
||
out.release()
|
||
print('total time: {:.4f}'.format(time.time() - tic_total))
|
||
#cv2.destroyAllWindows()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
args = set_retinaface_conf()
|
||
print(args.cpu)
|