|
@@ -0,0 +1,207 @@
|
|
|
+
|
|
|
+import os
|
|
|
+from utils import *
|
|
|
+from EmbedNet import *
|
|
|
+import torchvision.transforms as transforms
|
|
|
+from detectors import S3FD
|
|
|
+import argparse
|
|
|
+
|
|
|
+def createParser():
|
|
|
+ parser = argparse.ArgumentParser(description = "FaceNet");
|
|
|
+
|
|
|
+ parser.add_argument('--config', type=str, default=None, help='Config YAML file');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--batch_size', type=int, default=200, help='Batch size, number of classes per batch');
|
|
|
+ parser.add_argument('--max_img_per_cls', type=int, default=500, help='Maximum number of images per class per epoch');
|
|
|
+ parser.add_argument('--nDataLoaderThread', type=int, default=5, help='Number of loader threads');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--test_interval', type=int, default=5, help='Test and save every [test_interval] epochs');
|
|
|
+ parser.add_argument('--max_epoch', type=int, default=100, help='Maximum number of epochs');
|
|
|
+ parser.add_argument('--trainfunc', type=str, default="softmax", help='Loss function');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--optimizer', type=str, default="adam", help='sgd or adam');
|
|
|
+ parser.add_argument('--scheduler', type=str, default="steplr", help='Learning rate scheduler');
|
|
|
+ parser.add_argument('--lr', type=float, default=0.001, help='Learning rate');
|
|
|
+ parser.add_argument("--lr_decay", type=float, default=0.90, help='Learning rate decay every [test_interval] epochs');
|
|
|
+ parser.add_argument('--weight_decay', type=float, default=0, help='Weight decay in the optimizer');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument("--hard_prob", type=float, default=0.5, help='Hard negative mining probability, otherwise random, only for some loss functions');
|
|
|
+ parser.add_argument("--hard_rank", type=int, default=10, help='Hard negative mining rank in the batch, only for some loss functions');
|
|
|
+ parser.add_argument('--margin', type=float, default=0.1, help='Loss margin, only for some loss functions');
|
|
|
+ parser.add_argument('--scale', type=float, default=30, help='Loss scale, only for some loss functions');
|
|
|
+ parser.add_argument('--nPerClass', type=int, default=1, help='Number of images per class per batch, only for metric learning based losses');
|
|
|
+ parser.add_argument('--nClasses', type=int, default=8700, help='Number of classes in the softmax layer, only for softmax-based losses');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--initial_model', type=str, default="./models/amsoft_model.model", help='Initial model weights');
|
|
|
+ parser.add_argument('--save_path', type=str, default="exps/exp1", help='Path for model and logs');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--train_path', type=str, default="data/vggface2", help='Absolute path to the train set');
|
|
|
+ parser.add_argument('--train_ext', type=str, default="jpg", help='Training files extension');
|
|
|
+ parser.add_argument('--test_path', type=str, default="data/test", help='Absolute path to the test set');
|
|
|
+ parser.add_argument('--test_list', type=str, default="data/test_list.csv", help='Evaluation list');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--model', type=str, default="ResNet18", help='Name of model definition');
|
|
|
+ parser.add_argument('--nOut', type=int, default=512, help='Embedding size in the last FC layer');
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--eval', dest='eval', action='store_true', help='Eval only')
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--server', dest='server', action='store_true', help='Server mode')
|
|
|
+ parser.add_argument('--feat_save_path', type=str, default='saved_feats', help='Absolute path to the feature')
|
|
|
+ parser.add_argument('--port', type=int, default=10000, help='Port for the server')
|
|
|
+
|
|
|
+
|
|
|
+ parser.add_argument('--mixedprec', dest='mixedprec', action='store_true', help='Enable mixed precision training')
|
|
|
+
|
|
|
+ args = parser.parse_args()
|
|
|
+
|
|
|
+ return args
|
|
|
+
|
|
|
+def loadParameters(model, path):
|
|
|
+ state = model.state_dict()
|
|
|
+ loaded_state = torch.load(path)
|
|
|
+
|
|
|
+ for name, param in loaded_state.items():
|
|
|
+ origname = name;
|
|
|
+ if name not in state:
|
|
|
+ if name not in state:
|
|
|
+ print("%s is not in the model."%origname);
|
|
|
+ continue;
|
|
|
+
|
|
|
+ if state[name].size() != loaded_state[origname].size():
|
|
|
+ print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, state[name].size(), loaded_state[origname].size()));
|
|
|
+ continue;
|
|
|
+
|
|
|
+ state[name].copy_(param);
|
|
|
+
|
|
|
+DET = S3FD(device='cuda')
|
|
|
+app = Flask(__name__)
|
|
|
+args = createParser()
|
|
|
+UNKNOWN_THRESHOLD = 0.5
|
|
|
+
|
|
|
+
|
|
|
+s = EmbedNet(**vars(args)).cuda()
|
|
|
+transform = transforms.Compose(
|
|
|
+ [transforms.ToTensor(),
|
|
|
+ transforms.Resize(256),
|
|
|
+ transforms.CenterCrop([224,224]),
|
|
|
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+loadParameters(s, args.initial_model)
|
|
|
+s.eval()
|
|
|
+
|
|
|
+
|
|
|
+@app.route('/query', methods=['POST'])
|
|
|
+def query():
|
|
|
+
|
|
|
+ data = pickle.loads(request.get_data())
|
|
|
+
|
|
|
+ image = data['img']
|
|
|
+ image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
+
|
|
|
+ bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])
|
|
|
+
|
|
|
+ if len(bboxes) != 1:
|
|
|
+ return "fail"
|
|
|
+
|
|
|
+ bsi = 100
|
|
|
+
|
|
|
+ sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
|
|
|
+ sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
|
|
|
+ ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)
|
|
|
+
|
|
|
+ image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))
|
|
|
+
|
|
|
+ face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
|
|
|
+ face = cv2.resize(face,(240,240))
|
|
|
+
|
|
|
+ im1 = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
|
|
|
+
|
|
|
+ inp1 = transform(im1).cuda()
|
|
|
+
|
|
|
+ com_feat = s(inp1).detach().cpu()
|
|
|
+
|
|
|
+ files = glob.glob(os.path.join(args.feat_save_path, '*.pt'))
|
|
|
+
|
|
|
+ max_score = 0
|
|
|
+ pname = 'none'
|
|
|
+ for file in files:
|
|
|
+
|
|
|
+ ref_feat = torch.load(file)
|
|
|
+
|
|
|
+ score = F.cosine_similarity(ref_feat, com_feat)
|
|
|
+ if(score>max_score) :
|
|
|
+ max_score = score.item()
|
|
|
+ pname = file.split('/')[1].split('.')[0]
|
|
|
+
|
|
|
+ print('{} {:.2f}'.format(file,score.item()))
|
|
|
+
|
|
|
+ if max_score < UNKNOWN_THRESHOLD:
|
|
|
+ max_score = 0
|
|
|
+ pname = "Unknown"
|
|
|
+
|
|
|
+ return {
|
|
|
+ "file":pname,
|
|
|
+ "score":max_score,
|
|
|
+ "x1":bboxes[0][0],
|
|
|
+ "y1":bboxes[0][1],
|
|
|
+ "x2":bboxes[0][2],
|
|
|
+ "y2":bboxes[0][3]
|
|
|
+ }
|
|
|
+
|
|
|
+@app.route('/enroll', methods=['POST'])
|
|
|
+def enroll():
|
|
|
+
|
|
|
+ data = pickle.loads(request.get_data())
|
|
|
+
|
|
|
+ iname = data['name']
|
|
|
+
|
|
|
+ image = data['img']
|
|
|
+ image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
+
|
|
|
+ bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])
|
|
|
+
|
|
|
+ bsi = 100
|
|
|
+
|
|
|
+ sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
|
|
|
+ sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
|
|
|
+ ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)
|
|
|
+
|
|
|
+ image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))
|
|
|
+
|
|
|
+ face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
|
|
|
+ face = cv2.resize(face,(240,240))
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+ if not(os.path.exists(args.feat_save_path)):
|
|
|
+ os.makedirs(args.feat_save_path)
|
|
|
+
|
|
|
+ cv2.imwrite(os.path.join(args.feat_save_path, '{}.jpg'.format(iname)),face)
|
|
|
+ im1 = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
|
|
|
+
|
|
|
+ inp1 = transform(im1).cuda()
|
|
|
+
|
|
|
+ ref_feat = s(inp1).detach().cpu()
|
|
|
+
|
|
|
+ torch.save(ref_feat, os.path.join(args.feat_save_path,'{}.pt'.format(iname)))
|
|
|
+
|
|
|
+ return "success"
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ app.run(host='0.0.0.0', debug=True, port=args.port, threaded=False)
|