示例#1
0
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):
示例#3
0
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')            
示例#4
0
                        '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