1 files changed, 56 insertions, 0 deletions
diff --git a/megapixels/app/processors/face_recognition.py b/megapixels/app/processors/face_recognition.py
new file mode 100644
index 00000000..e0b9f752
--- /dev/null
+++ b/megapixels/app/processors/face_recognition.py
@@ -0,0 +1,56 @@
+import os
+from os.path import join
+from pathlib import Path
+
+import cv2 as cv
+import numpy as np
+import dlib
+import imutils
+
+from app.utils import im_utils, logger_utils
+from app.models.bbox import BBox
+from app.settings import app_cfg as cfg
+from app.settings import types
+    
+class RecognitionDLIB:
+
+  # https://github.com/davisking/dlib/blob/master/python_examples/face_recognition.py
+  # facerec.compute_face_descriptor(img, shape, 100, 0.25)
+
+  def __init__(self, gpu=0):
+    self.log = logger_utils.Logger.getLogger()
+    
+    if gpu > -1:
+      cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '')
+      os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
+    
+    self.predictor = dlib.shape_predictor(cfg.DIR_MODELS_DLIB_5PT)
+    self.facerec = dlib.face_recognition_model_v1(cfg.DIR_MODELS_DLIB_FACEREC_RESNET)
+    
+    if gpu > -1:
+      os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices  # reset GPU env
+
+
+  def vec(self, im, bbox, width=100, 
+      jitters=cfg.DLIB_FACEREC_JITTERS, padding=cfg.DLIB_FACEREC_PADDING):
+    # Converts image and bbox into 128d vector
+    # scale the image so the face is always 100x100 pixels
+
+    #self.log.debug('compute scale')
+    scale = width / bbox.width
+    #im = cv.resize(im, (scale, scale), cv.INTER_LANCZOS4)
+    #self.log.debug('resize')
+    cv.resize(im, None, fx=scale, fy=scale, interpolation=cv.INTER_LANCZOS4) 
+    #self.log.debug('to dlib')
+    bbox_dlib = bbox.to_dlib()
+    #self.log.debug('precitor')
+    face_shape = self.predictor(im, bbox_dlib)
+    # vec = self.facerec.compute_face_descriptor(im, face_shape, jitters, padding)
+    #self.log.debug('vec')
+    vec = self.facerec.compute_face_descriptor(im, face_shape, jitters)
+    #vec = self.facerec.compute_face_descriptor(im, face_shape)
+    return vec
+
+
+  def similarity(self, query_enc, known_enc):
+    return np.linalg.norm(query_enc - known_enc, axis=1)