1 files changed, 43 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..9c3a301d
--- /dev/null
+++ b/megapixels/app/processors/face_recognition.py
@@ -0,0 +1,43 @@
+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, opt_gpu=0):
+    self.log = logger_utils.Logger.getLogger()
+    if opt_gpu > 0:
+      cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '')
+      os.environ['CUDA_VISIBLE_DEVICES'] = str(opt_gpu)
+      self.predictor = dlib.shape_predictor(cfg.DIR_MODELS_DLIB_5PT)
+      self.facerec = dlib.face_recognition_model_v1(cfg.DIR_MODELS_DLIB_FACEREC_RESNET)
+      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
+
+    scale = width / bbox.width
+    im = cv.resize(im, (scale, scale), interploation=cv.INTER_LANCZOS4)
+    bbox_dlib = bbox.to_dlib()
+    face_shape = self.predictor(im, bbox_dlib)
+    vec = self.facerec.compute_face_descriptor(im, face_shape, jitters, padding)
+    return vec
+
+
+  def similarity(self, query_enc, known_enc):
+    return np.linalg.norm(query_enc - known_enc, axis=1)