def get_detect_model():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
detect_model = MobileFaceNet(512).to(
device) # embeding size is 512 (feature vector)
detect_model.load_state_dict(
torch.load('Weights/MobileFace_Net.pt',
map_location=lambda storage, loc: storage))
print('MobileFaceNet face detection model generated')
detect_model.eval()
return detect_model
import argparse
import sys
import os
import io
import json
from torchvision import models
from PIL import Image
from flask import Flask, jsonify, request
from flask_cors import CORS
device_0 = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
detect_model = MobileFaceNet(512).to(device_0)
detect_model.load_state_dict(
torch.load('Weights/MobileFace_Net',
map_location=lambda storage, loc: storage))
detect_model.eval()
target, name = load_facebank(path='facebank')
parser = argparse.ArgumentParser()
parser.add_argument('--miniface', default=10, type=int)
parser.add_argument('--scale', default=2, type=int)
parser.add_argument('--update', default=False, type=bool)
args = parser.parse_args()
if args.update:
targets, names = prepare_facebank(detect_model, path='facebank')
print('facebank updated')
else:
targets, names = load_facebank(path='facebank')
print('facebank loaded')
def mod_crop(image, scale=2):
class VideoCamera(object):
def __init__(self,username='******',threshold=80,update_val=True,tta=False,c=True,scale=0.3,min_face=20,embedding=512,bank_path='media/account/facebank',model_path='Weights/model_final.pth'):
self.video = cv2.VideoCapture(0)
(self.grabbed, self.frame) = self.video.read()
self.thread=threading.Thread(target=self.update, args=())
self.flag=True
# recognition parameter
self.threshold=threshold
self.tta=tta
self.score=c
self.scale=scale
self.min_face=min_face
self.embedding=embedding
self.facebank='media/account/{}/facebank'.format(username)
self.model_path=model_path
self.up=update_val
self.device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('using ',self.device)
self.face_detector=mtcnn_custom(device=self.device,p_model_path='MTCNN/weights/pnet_Weights',r_model_path='MTCNN/weights/rnet_Weights',o_model_path='MTCNN/weights/onet_Weights')
print('face detector created...')
#prepare pretrained model
self.detect_model = MobileFaceNet(self.embedding).to(self.device) # embeding size is 512 (feature vector)
# self.check_point=torch.load('Weights/model_final_t.pth',map_location=lambda storage, loc: storage)
# self.detect_model.load_state_dict(self.check_point['model_state_dict'])
self.detect_model.load_state_dict(torch.load('Weights/model_final.pth',map_location=lambda storage, loc: storage))
print('MobileFaceNet face detection model generated')
self.detect_model.eval()
#face bank update
if self.up:
self.targets, self.names = prepare_facebank(self.detect_model, path=self.facebank, tta=self.tta)
print('facebank updated')
else:
self.targets, self.names = load_facebank(path=self.facebank)
print('facebank loaded')
# targets: number of candidate x 512
def __del__(self):
self.video.release()
def update_facebank(self,img_list):
for img_path in img_list:
img=cv2.imread(img_path)
bboxes,landmarks=self.face_detector.detect_all_net(image=img,mini_face=self.min_face)
faces= Face_alignment(img,default_square=True,landmarks=landmarks)
try:
os.remove(img_path)
except:
print('fail to remove')
cv2.imwrite(img_path,faces[0])
self.targets,self.names=prepare_facebank(self.detect_model, path=self.facebank, tta=self.tta)
print('new facebank uploaded !!...')
def get_frame(self): #여기다가 face detection, recognition기능을 넣으면 문제없음.
frame = self.frame #thread가 update하는 이미지를 가져 옴.
if frame is not None and self.flag is True:
try:
start_time=time.time()
input=resize_image(frame,self.scale)# input size를 줄여줌으로 speed up 가능
#print('get bboxes')
# bboxes, landmarks = create_mtcnn_net(input, self.min_face, self.device, p_model_path='MTCNN/weights/pnet_Weights',
# r_model_path='MTCNN/weights/rnet_Weights',
# o_model_path='MTCNN/weights/onet_Weights')
#print('sucess bbox')
bboxes,landmarks=self.face_detector.detect_all_net(image=input,mini_face=self.min_face)
if bboxes != []:
bboxes=bboxes/self.scale
landmarks=landmarks/self.scale
faces= Face_alignment(frame,default_square=True,landmarks=landmarks)
embs=[]
test_transform = trans.Compose([
trans.ToTensor(),
trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
for img in faces:
if self.tta:
mirror = cv2.flip(img,1)
emb = self.detect_model(test_transform(img).to(self.device).unsqueeze(0))
emb_mirror = self.detect_model(test_transform(mirror).to(self.device).unsqueeze(0))
embs.append(l2_norm(emb + emb_mirror))
else:
embs.append(self.detect_model(test_transform(img).to(self.device).unsqueeze(0)))
source_embs=torch.cat(embs)
diff=source_embs.unsqueeze(-1) - self.targets.transpose(1, 0).unsqueeze(0) # i.e
dist = torch.sum(torch.pow(diff, 2), dim=1) # number of detected faces x numer of target faces
minimum, min_idx = torch.min(dist, dim=1) # min and idx for each row
min_idx[minimum > ((self.threshold-156)/(-80))] = -1 # if no match, set idx to -1
score = minimum
results = min_idx
score_100 = torch.clamp(score*-80+156,0,100)
FPS=1.0/(time.time()-start_time)
cv2.putText(frame,'FPS : {:.1f}'.format(FPS),(10,15),cv2.FONT_HERSHEY_DUPLEX,0.75,(255,0,255))
for i,b in enumerate(bboxes):
b=b.astype('uint32')
cv2.rectangle(frame,(b[0],b[1]),(b[2],b[3]),(0,255,0),1)
try:
if self.names[results[i]+1]=='Unknown': #mosic func
#print('detect unknwon')
face_region=frame[b[1]:b[3],b[0]:b[2]]
face_region=cv2.blur(face_region,(30,30))
frame[b[1]:b[3],b[0]:b[2]]=face_region
except:
pass
# development version
# if self.score:
# cv2.putText(frame,self.names[results[i]+1]+' score:{:.0f}'.format(score_100[i]),(int(b[0]),int(b[1]-25)),cv2.FONT_ITALIC,1,(255,255,0))
# else:
# cv2.putText(frame,self.names[results[i]+1],(int(b[0]),int(b[1]-25)),cv2.FONT_ITALIC,1,(255,255,0))
except:
pass
_, jpeg = cv2.imencode('.jpg',frame)
return jpeg.tobytes()
def update(self): # thread를 사용하여 지속적으로 cap.read반복
while True:
(self.grabbed, self.frame) = self.video.read()
if self.grabbed==False or self.flag==False:
break
def gen(self):
while True:
if self.flag==False:
break
frame = self.get_frame()
if frame is None:
continue
yield(b'--frame\r\n'b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')
'net_state_dict': model.state_dict()
},
os.path.join(save_dir,
'Iter_%06d_model.ckpt' % total_iters))
torch.save(
{
'iters': total_iters,
'net_state_dict': margin.state_dict()
},
os.path.join(save_dir,
'Iter_%06d_margin.ckpt' % total_iters))
# evaluate accuracy
if total_iters % 3000 == 0:
model.eval()
for phase in ['LFW', 'CFP_FP', 'AgeDB30']:
featureLs, featureRs = getFeature(model,
dataloaders[phase],
device,
flip=args.flip)
ACCs, threshold = evaluation_10_fold(featureLs,
featureRs,
dataset[phase],
method=args.method)
print(
'Epoch {}/{}，{} average acc:{:.4f} average threshold:{:.4f}'
.format(epoch, args.epoch - 1, phase,
np.mean(ACCs) * 100, np.mean(threshold)))
if best_acc[phase] <= np.mean(ACCs) * 100:
best_acc[phase] = np.mean(ACCs) * 100