133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
from datetime import datetime
|
|
import numpy as np
|
|
import io, cv2, os
|
|
from .model import l2_norm
|
|
import torch
|
|
import matplotlib.pyplot as plt
|
|
plt.switch_backend('agg')
|
|
|
|
|
|
def faces_preprocessing(faces, device):
|
|
faces = faces.permute(0, 3, 1, 2).float()
|
|
faces = faces.div(255).to(device)
|
|
mu = torch.as_tensor([.5, .5, .5], dtype=faces.dtype, device=device)
|
|
faces[:].sub_(mu[:, None, None]).div_(mu[:, None, None])
|
|
return faces
|
|
|
|
|
|
def separate_bn_paras(modules):
|
|
if not isinstance(modules, list):
|
|
modules = [*modules.modules()]
|
|
paras_only_bn = []
|
|
paras_wo_bn = []
|
|
for layer in modules:
|
|
if 'model' in str(layer.__class__):
|
|
continue
|
|
if 'container' in str(layer.__class__):
|
|
continue
|
|
else:
|
|
if 'batchnorm' in str(layer.__class__):
|
|
paras_only_bn.extend([*layer.parameters()])
|
|
else:
|
|
paras_wo_bn.extend([*layer.parameters()])
|
|
return paras_only_bn, paras_wo_bn
|
|
|
|
|
|
def update_facebank(conf, model, detector, tta=True):
|
|
if os.path.exists(conf.facebank_path) == False:
|
|
raise Exception("you don't have facebank yet: create with add_from_webcam or add_from_folder function")
|
|
model.eval()
|
|
faces_embs = torch.empty(0).to(conf.device)
|
|
names = np.array(['Unknown'])
|
|
for path in conf.facebank_path.iterdir():
|
|
if path.is_file():
|
|
continue
|
|
faces = []
|
|
for img_path in path.iterdir():
|
|
face = cv2.imread(str(img_path))
|
|
if face.shape[:2] != (112, 112): # if img be not face
|
|
face = detector.detect_align(face)[0]
|
|
cv2.imwrite(img_path, face)
|
|
else:
|
|
face = torch.tensor(face).unsqueeze(0)
|
|
faces.append(face)
|
|
|
|
faces = torch.cat(faces)
|
|
if len(faces.shape) <= 3:
|
|
continue
|
|
|
|
with torch.no_grad():
|
|
faces = faces_preprocessing(faces, device=conf.device)
|
|
if tta:
|
|
face_emb = model(faces)
|
|
hflip_emb = model(faces.flip(-1)) # image horizontal flip
|
|
face_embs = l2_norm(face_emb + hflip_emb)
|
|
else:
|
|
face_embs = model(faces)
|
|
|
|
faces_embs = torch.cat((faces_embs, face_embs.mean(0, keepdim=True)))
|
|
names = np.append(names, path.name)
|
|
|
|
torch.save(faces_embs, conf.facebank_path/'facebank.pth')
|
|
np.save(conf.facebank_path/'names', names)
|
|
print('from FaceRecognizer: facebank updated')
|
|
return faces_embs, names
|
|
|
|
|
|
def load_facebank(conf):
|
|
if os.path.exists(conf.facebank_path) == False:
|
|
raise Exception("you don't have facebank yet: create with add_from_webcam or add_from_folder function")
|
|
|
|
embs = torch.load(conf.facebank_path/'facebank.pth')
|
|
names = np.load(conf.facebank_path/'names.npy')
|
|
print('from FaceRecognizer: facebank loaded')
|
|
return embs, names
|
|
|
|
|
|
def get_time():
|
|
return (str(datetime.now())[:-10]).replace(' ', '-').replace(':', '-')
|
|
|
|
|
|
def gen_plot(fpr, tpr):
|
|
"""Create a pyplot plot and save to buffer."""
|
|
plt.figure()
|
|
plt.xlabel("FPR", fontsize=14)
|
|
plt.ylabel("TPR", fontsize=14)
|
|
plt.title("ROC Curve", fontsize=14)
|
|
plt.plot(fpr, tpr, linewidth=2)
|
|
buf = io.BytesIO()
|
|
plt.savefig(buf, format='jpeg')
|
|
buf.seek(0)
|
|
plt.close()
|
|
return buf
|
|
|
|
|
|
def draw(bbox, name, frame):
|
|
frame = cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 6)
|
|
frame = cv2.putText(frame, name, (bbox[0], bbox[1]), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 4, cv2.LINE_AA)
|
|
return frame
|
|
|
|
|
|
def special_draw(img, box, landmarsk, name, score=100):
|
|
"""draw a bounding box on image"""
|
|
color = (148, 133, 0)
|
|
tl = round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line thickness
|
|
c1, c2 = (box[0], box[1]), (box[2], box[3])
|
|
# draw bounding box
|
|
cv2.rectangle(img, c1, c2, color, thickness=tl)
|
|
# draw landmark
|
|
for land in landmarsk:
|
|
cv2.circle(img, tuple(land.int().tolist()), 3, color, -1)
|
|
# draw score
|
|
score = 100-(score*100/1.4)
|
|
score = 0 if score < 0 else score
|
|
bar = (box[3] + 2) - (box[1] - 2)
|
|
score_final = bar - (score*bar/100)
|
|
cv2.rectangle(img, (box[2] + 1, box[1] - 2 + score_final), (box[2] + (tl+5), box[3] + 2), color, -1)
|
|
# draw label
|
|
tf = max(tl - 1, 1) # font thickness
|
|
t_size = cv2.getTextSize(name, 0, fontScale=tl / 3, thickness=tf)[0]
|
|
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
|
|
cv2.rectangle(img, c1, c2, color, -1) # filled
|
|
cv2.putText(img, name, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|