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