server.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272
  1. import os, datetime, numpy as np
  2. from utils import *
  3. from EmbedNet import *
  4. import torchvision.transforms as transforms
  5. from detectors import S3FD
  6. import argparse
  7. def createParser():
  8. parser = argparse.ArgumentParser(description = "FaceNet");
  9. parser.add_argument('--config', type=str, default=None, help='Config YAML file');
  10. ## Data loader
  11. parser.add_argument('--batch_size', type=int, default=200, help='Batch size, number of classes per batch');
  12. parser.add_argument('--max_img_per_cls', type=int, default=500, help='Maximum number of images per class per epoch');
  13. parser.add_argument('--nDataLoaderThread', type=int, default=5, help='Number of loader threads');
  14. ## Training details
  15. parser.add_argument('--test_interval', type=int, default=5, help='Test and save every [test_interval] epochs');
  16. parser.add_argument('--max_epoch', type=int, default=100, help='Maximum number of epochs');
  17. parser.add_argument('--trainfunc', type=str, default="softmax", help='Loss function');
  18. ## Optimizer
  19. parser.add_argument('--optimizer', type=str, default="adam", help='sgd or adam');
  20. parser.add_argument('--scheduler', type=str, default="steplr", help='Learning rate scheduler');
  21. parser.add_argument('--lr', type=float, default=0.001, help='Learning rate');
  22. parser.add_argument("--lr_decay", type=float, default=0.90, help='Learning rate decay every [test_interval] epochs');
  23. parser.add_argument('--weight_decay', type=float, default=0, help='Weight decay in the optimizer');
  24. ## Loss functions
  25. parser.add_argument("--hard_prob", type=float, default=0.5, help='Hard negative mining probability, otherwise random, only for some loss functions');
  26. parser.add_argument("--hard_rank", type=int, default=10, help='Hard negative mining rank in the batch, only for some loss functions');
  27. parser.add_argument('--margin', type=float, default=0.1, help='Loss margin, only for some loss functions');
  28. parser.add_argument('--scale', type=float, default=30, help='Loss scale, only for some loss functions');
  29. parser.add_argument('--nPerClass', type=int, default=1, help='Number of images per class per batch, only for metric learning based losses');
  30. parser.add_argument('--nClasses', type=int, default=8700, help='Number of classes in the softmax layer, only for softmax-based losses');
  31. ## Load and save
  32. parser.add_argument('--initial_model', type=str, default="./models/amsoft_model.model", help='Initial model weights');
  33. parser.add_argument('--save_path', type=str, default="exps/exp1", help='Path for model and logs');
  34. ## Training and test data
  35. parser.add_argument('--train_path', type=str, default="data/vggface2", help='Absolute path to the train set');
  36. parser.add_argument('--train_ext', type=str, default="jpg", help='Training files extension');
  37. parser.add_argument('--test_path', type=str, default="data/test", help='Absolute path to the test set');
  38. parser.add_argument('--test_list', type=str, default="data/test_list.csv", help='Evaluation list');
  39. ## Model definition
  40. parser.add_argument('--model', type=str, default="ResNet18", help='Name of model definition');
  41. parser.add_argument('--nOut', type=int, default=512, help='Embedding size in the last FC layer');
  42. ## For test only
  43. parser.add_argument('--eval', dest='eval', action='store_true', help='Eval only')
  44. ## For server
  45. parser.add_argument('--server', dest='server', action='store_true', help='Server mode')
  46. parser.add_argument('--feat_save_path', type=str, default='saved_feats', help='Absolute path to the feature')
  47. parser.add_argument('--img_save_path', type=str, default='saved_img', help='Absolute path to the image')
  48. parser.add_argument('--port', type=int, default=10000, help='Port for the server')
  49. ## Distributed and mixed precision training
  50. parser.add_argument('--mixedprec', dest='mixedprec', action='store_true', help='Enable mixed precision training')
  51. args = parser.parse_args()
  52. return args
  53. def loadParameters(model, path):
  54. state = model.state_dict()
  55. loaded_state = torch.load(path)
  56. for name, param in loaded_state.items():
  57. origname = name;
  58. if name not in state:
  59. if name not in state:
  60. print("%s is not in the model."%origname);
  61. continue;
  62. if state[name].size() != loaded_state[origname].size():
  63. print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, state[name].size(), loaded_state[origname].size()));
  64. continue;
  65. state[name].copy_(param);
  66. class your_dataset(torch.utils.data.Dataset):
  67. def __init__(self, files):
  68. self.data = files
  69. print('{:d} files in the dataset'.format(len(self.data)))
  70. def __getitem__(self, index):
  71. fname = self.data[index]
  72. try:
  73. # return image if read is successful
  74. image = cv2.imread(fname)
  75. image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  76. return image, image_np, fname
  77. except:
  78. # return empty if not successful
  79. return np.array([]), np.array([]), fname
  80. def __len__(self):
  81. return len(self.data)
  82. DET = S3FD(device='cuda')
  83. app = Flask(__name__)
  84. args = createParser()
  85. UNKNOWN_THRESHOLD = 0.5
  86. s = EmbedNet(**vars(args)).cuda()
  87. transform = transforms.Compose(
  88. [transforms.ToTensor(),
  89. transforms.Resize(256),
  90. transforms.CenterCrop([224,224]),
  91. transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
  92. # trainer = ModelTrainer(s, **vars(args))
  93. loadParameters(s, args.initial_model)
  94. s.eval()
  95. @app.route('/cal', methods=['POST'])
  96. def calculate():
  97. # unpack the received data
  98. data = pickle.loads(request.get_data())
  99. iname = data['name']
  100. image_save_path = os.path.join(args.img_save_path, iname)
  101. files = glob.glob(os.path.join(image_save_path, '{}*.jpg'.format(iname)))
  102. dataset = your_dataset(files)
  103. loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=10)
  104. embedding_list = ''
  105. for data in loader:
  106. image = data[0][0].numpy()
  107. image_np = data[1][0].numpy()
  108. fname = data[2][0].split('/')[1]
  109. bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])
  110. bsi = 100
  111. sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
  112. sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
  113. ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)
  114. image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))
  115. face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
  116. face = cv2.resize(face,(240,240))
  117. im1 = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
  118. inp1 = transform(im1).cuda()
  119. ref_feat = s(inp1).detach().cpu()
  120. embedding_list = torch.cat([ref_feat, embedding_list]) if embedding_list != '' else ref_feat
  121. embedding_mean = torch.mean(embedding_list,dim=0,keepdim=True)
  122. torch.save(embedding_mean, os.path.join(args.feat_save_path,'{}.pt'.format(fname)))
  123. return 'success'
  124. @app.route('/query', methods=['POST'])
  125. def query():
  126. # unpack the received data
  127. data = pickle.loads(request.get_data())
  128. image = data['img']
  129. image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  130. bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])
  131. if len(bboxes) != 1:
  132. return "fail"
  133. bsi = 100
  134. sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
  135. sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
  136. ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)
  137. image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))
  138. face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
  139. face = cv2.resize(face,(240,240))
  140. im1 = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
  141. inp1 = transform(im1).cuda()
  142. com_feat = s(inp1).detach().cpu()
  143. files = glob.glob(os.path.join(args.feat_save_path, '*.pt'))
  144. max_score = 0
  145. pname = 'none'
  146. for file in files:
  147. ref_feat = torch.load(file)
  148. score = F.cosine_similarity(ref_feat, com_feat)
  149. if(score>max_score) :
  150. max_score = score.item()
  151. pname = file.split('/')[1].split('.')[0]
  152. print('{} {:.2f}'.format(file,score.item()))
  153. if max_score < 0.1:
  154. max_score = 0
  155. pname = "Unknown"
  156. return {
  157. "file":pname,
  158. "score":max_score,
  159. "x1":bboxes[0][0],
  160. "y1":bboxes[0][1],
  161. "x2":bboxes[0][2],
  162. "y2":bboxes[0][3]
  163. }
  164. @app.route('/enroll', methods=['POST'])
  165. def enroll():
  166. # unpack the received data
  167. data = pickle.loads(request.get_data())
  168. iname = data['name']
  169. image = data['img']
  170. image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  171. bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])
  172. if len(bboxes) != 1:
  173. return "fail"
  174. bsi = 100
  175. sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
  176. sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
  177. ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)
  178. image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))
  179. face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
  180. face = cv2.resize(face,(240,240))
  181. now = datetime.datetime.now().strftime('%y-%m-%d-%H-%M-%f')
  182. image_save_path = os.path.join(args.img_save_path, iname)
  183. if not(os.path.exists(args.feat_save_path)):
  184. os.makedirs(args.feat_save_path)
  185. if not(os.path.exists(image_save_path)):
  186. os.makedirs(image_save_path)
  187. cv2.imwrite(os.path.join(image_save_path, '{}_{}.jpg'.format(iname, now)), face)
  188. return "success"
  189. if __name__ == "__main__":
  190. app.run(host='0.0.0.0', debug=True, port=args.port, threaded=False)