Making face recognition system independent of dataset. Major refactor…

…ing.

Former-commit-id: ebeed61

__init__.py

@@ -1,4 +1,4 @@
-from .arsfutura_face_recognition import face_recogniser_factory
+from .face_recognition import face_recogniser_factory
 
 face_recognizer = face_recogniser_factory(include_predictions=True)
 

face_recognition.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-from arsfutura_face_recognition import face_recogniser_factory
+import joblib
 import argparse
 import cv2
 import numpy as np
@@ -11,6 +11,7 @@ def parse_args(args=None):
     parser = argparse.ArgumentParser(
         'Script for recognising faces on picture. Output of this script is json with list of people on picture and '
         'base64 encoded picture which has bounding boxes of people.')
+    parser.add_argument('-m', '--model-path', required=True, help='Path to face recogniser model.')
     parser.add_argument('--image-path', required=True, help='Path to image file.')
     return parser.parse_args(args)
 
@@ -25,8 +26,8 @@ def draw_bb_on_img(faces, img):
 
 
 def _recognise_faces(args):
-    img = Image.open(args.image_path)
-    faces = face_recogniser_factory(include_predictions=True)(img)
+    img = Image.open(args.image_path).convert('RGB')
+    faces = joblib.load(args.model_path)(img)
     img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
     draw_bb_on_img(faces, img_cv)
     return faces, img_cv

face_recognition/__init__.py

@@ -0,0 +1,2 @@
+from .face_recogniser import FaceRecogniser
+from .face_features_extractor import FaceFeaturesExtractor

face_recognition/face_features_extractor.py

@@ -0,0 +1,27 @@
+import torch
+from facenet_pytorch import MTCNN, InceptionResnetV1
+from torchvision import transforms
+from . import preprocessing
+from facenet_pytorch.models.utils.detect_face import extract_face
+
+
+class FaceFeaturesExtractor:
+    def __init__(self):
+        self.aligner_preprocess = transforms.Compose([preprocessing.ExifOrientationNormalize()])
+        self.aligner = MTCNN(prewhiten=False, keep_all=True)
+        self.facenet_preprocess = transforms.Compose([preprocessing.Whitening()])
+        self.facenet = InceptionResnetV1(pretrained='vggface2').eval()
+
+    def extract_features(self, img):
+        bbs, _ = self.aligner.detect(self.aligner_preprocess(img))
+        if bbs is None:
+            # if no face is detected
+            return None, None
+
+        faces = torch.stack([extract_face(img, bb) for bb in bbs])
+        embeddings = self.facenet(faces).detach().numpy()
+
+        return bbs, embeddings
+
+    def __call__(self, img):
+        return self.extract_features(img)

face_recognition/face_recogniser.py

@@ -0,0 +1,41 @@
+from collections import namedtuple
+
+Prediction = namedtuple('Prediction', 'id name confidence')
+Face = namedtuple('Face', 'top_prediction bb all_predictions')
+BoundingBox = namedtuple('BoundingBox', 'left top right bottom')
+
+
+def top_prediction(idx_to_class, probs):
+    top_label = probs.argmax()
+    return Prediction(id=top_label, name=idx_to_class[top_label], confidence=probs[top_label])
+
+
+def to_predictions(idx_to_class, probs):
+    return [Prediction(id=i, name=idx_to_class[i], confidence=prob) for i, prob in enumerate(probs)]
+
+
+class FaceRecogniser:
+    def __init__(self, feature_extractor, classifier, idx_to_class):
+        self.feature_extractor = feature_extractor
+        self.classifier = classifier
+        self.idx_to_class = idx_to_class
+
+    def recognise_faces(self, img):
+        bbs, embeddings = self.feature_extractor(img)
+        if bbs is None:
+            # if no faces are detected
+            return []
+
+        predictions = self.classifier.predict_proba(embeddings)
+
+        return [
+            Face(
+                top_prediction=top_prediction(self.idx_to_class, probs),
+                bb=BoundingBox(left=bb[0], top=bb[1], right=bb[2], bottom=bb[3]),
+                all_predictions=to_predictions(self.idx_to_class, probs)
+            )
+            for bb, probs in zip(bbs, predictions)
+        ]
+
+    def __call__(self, img):
+        return self.recognise_faces(img)

real_time_face_detection.py

@@ -1,13 +1,21 @@
 #!/usr/bin/env python
 
+import argparse
+import joblib
 import cv2
 from PIL import Image
-from arsfutura_face_recognition import face_recogniser_factory
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-m', '--model-path', required=True, help='Path to face recogniser model.')
+    return parser.parse_args()
 
 
 def main():
+    args = parse_args()
     cap = cv2.VideoCapture(0)
-    face_recogniser = face_recogniser_factory()
+    face_recogniser = joblib.load(args.model_path)
 
     while True:
         # Capture frame-by-frame