testing on hugging

models/mobileNet.tflite

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6c19b789f661caa8da735566490bfd8895beffb2a1ec97a56b126f0539991aa6
-size 8210384

pyproject.toml

@@ -0,0 +1,3 @@
+[tool.black]
+line-length = 120
+target-version = ['py38']

requirements.txt

@@ -8,4 +8,6 @@ streamlit-webrtc
 matplotlib
 streamlit-toggle-switch
 tflite-runtime
-twilio
+twilio
+tqdm
+plotly

tools/annotation.py

@@ -0,0 +1,107 @@
+import numpy as np
+import cv2
+
+
+class Annotation:
+    def __init__(self, draw_bbox=True, draw_landmarks=True, draw_name=True, upscale=True):
+        self.bbox = draw_bbox
+        self.landmarks = draw_landmarks
+        self.name = draw_name
+        self.upscale = upscale
+
+    def __call__(self, frame, detections, identities, matches, gallery):
+        shape = np.asarray(frame.shape[:2][::-1])
+        if self.upscale:
+            frame = cv2.resize(frame, (1920, 1080))
+            upscale_factor = np.asarray([1920 / shape[0], 1080 / shape[1]])
+            shape = np.asarray(frame.shape[:2][::-1])
+        else:
+            upscale_factor = np.asarray([1, 1])
+
+        frame.flags.writeable = True
+
+        for detection in detections:
+            # Draw Landmarks
+            if self.landmarks:
+                for landmark in detection.landmarks:
+                    cv2.circle(
+                        frame,
+                        (landmark * upscale_factor).astype(int),
+                        2,
+                        (255, 255, 255),
+                        -1,
+                    )
+
+            # Draw Bounding Box
+            if self.bbox:
+                cv2.rectangle(
+                    frame,
+                    (detection.bbox[0] * upscale_factor).astype(int),
+                    (detection.bbox[1] * upscale_factor).astype(int),
+                    (255, 0, 0),
+                    2,
+                )
+
+            # Draw Index
+            cv2.putText(
+                frame,
+                str(detection.idx),
+                (
+                    ((detection.bbox[1][0] + 2) * upscale_factor[0]).astype(int),
+                    ((detection.bbox[1][1] + 2) * upscale_factor[1]).astype(int),
+                ),
+                cv2.LINE_AA,
+                0.5,
+                (0, 0, 0),
+                2,
+            )
+
+        # Draw Name
+        if self.name:
+            for match in matches:
+                try:
+                    detection = detections[identities[match.identity_idx].detection_idx]
+                except:
+                    print("Identity IDX: ", match.identity_idx)
+                    print("Len(Detections): ", len(detections))
+                    print("Len(Identites): ", len(identities))
+                    print("Detection IDX: ", identities[match.identity_idx].detection_idx)
+
+                    # print("Detections: ", detections)
+
+                cv2.rectangle(
+                    frame,
+                    (detection.bbox[0] * upscale_factor).astype(int),
+                    (detection.bbox[1] * upscale_factor).astype(int),
+                    (0, 255, 0),
+                    2,
+                )
+
+                cv2.rectangle(
+                    frame,
+                    (
+                        (detection.bbox[0][0] * upscale_factor[0]).astype(int),
+                        (detection.bbox[0][1] * upscale_factor[1] - (shape[1] // 25)).astype(int),
+                    ),
+                    (
+                        (detection.bbox[1][0] * upscale_factor[0]).astype(int),
+                        (detection.bbox[0][1] * upscale_factor[1]).astype(int),
+                    ),
+                    (255, 255, 255),
+                    -1,
+                )
+
+                cv2.putText(
+                    frame,
+                    gallery[match.gallery_idx].name,
+                    (
+                        ((detection.bbox[0][0] + shape[0] // 400) * upscale_factor[0]).astype(int),
+                        ((detection.bbox[0][1] - shape[1] // 100) * upscale_factor[1]).astype(int),
+                    ),
+                    cv2.LINE_AA,
+                    0.5,
+                    (0, 0, 0),
+                    2,
+                )
+
+        return frame

tools/face_detection.py

@@ -0,0 +1,481 @@
+import tflite_runtime.interpreter as tflite
+import cv2
+import numpy as np
+from .utils import tflite_inference
+from .nametypes import Detection
+from .utils import get_file
+
+
+BASE_URL = "https://github.com/Martlgap/FaceIDLight/releases/download/v.0.1/"
+
+FILE_HASHES = {
+    "o_net": "768385d570300648b7b881acbd418146522b79b4771029bb2e684bdd8c764b9f",
+    "p_net": "530183192e24f7cc86b6706e1eb600482c4ed4306399ac939c472e3957bae15e",
+    "r_net": "5ec33b065eb2802bc4c2575d21feff1a56958d854785bc3e2907d3b7ace861a2",
+}
+
+
+class StageStatus:
+    """
+    Keeps status between MTCNN stages
+    """
+
+    def __init__(self, pad_result: tuple = None, width=0, height=0):
+        self.width = width
+        self.height = height
+        self.dy = self.edy = self.dx = self.edx = self.y = self.ey = self.x = self.ex = self.tmp_w = self.tmp_h = []
+
+        if pad_result is not None:
+            self.update(pad_result)
+
+    def update(self, pad_result: tuple):
+        s = self
+        s.dy, s.edy, s.dx, s.edx, s.y, s.ey, s.x, s.ex, s.tmp_w, s.tmp_h = pad_result
+
+
+class FaceDetection:
+    """
+    Allows to perform MTCNN Detection ->
+        a) Detection of faces (with the confidence probability)
+        b) Detection of keypoints (left eye, right eye, nose, mouth_left, mouth_right)
+    """
+
+    def __init__(
+        self,
+        min_face_size: int = 40,
+        steps_threshold: list = None,
+        scale_factor: float = 0.7,
+        min_detections_conf: float = 0.9,
+    ):
+        """
+        Initializes the MTCNN.
+        :param min_face_size: minimum size of the face to detect
+        :param steps_threshold: step's thresholds values
+        :param scale_factor: scale factor
+        """
+        if steps_threshold is None:
+            steps_threshold = [0.6, 0.7, 0.7]  # original mtcnn values [0.6, 0.7, 0.7]
+        self._min_face_size = min_face_size
+        self._steps_threshold = steps_threshold
+        self._scale_factor = scale_factor
+        self.min_detections_conf = min_detections_conf
+        self.p_net = tflite.Interpreter(model_path=get_file(BASE_URL + "p_net.tflite", FILE_HASHES["p_net"]))
+        self.r_net = tflite.Interpreter(model_path=get_file(BASE_URL + "r_net.tflite", FILE_HASHES["r_net"]))
+        self.o_net = tflite.Interpreter(model_path=get_file(BASE_URL + "o_net.tflite", FILE_HASHES["o_net"]))
+
+    def __call__(self, frame):
+        """
+        Detects bounding boxes from the specified image.
+        :param img: image to process
+        :return: list containing all the bounding boxes detected with their keypoints.
+
+        From MTCNN:
+        # Total boxes (bBoxes for faces)
+        # 1. dim -> Number of found Faces
+        # 2. dim -> x_min, y_min, x_max, y_max, score
+
+        # Points (Landmarks left eye, right eye, nose, left mouth, right mouth)
+        # 1. dim -> Number of found Faces
+        # 2. dim -> x1, x2, x3, x4, x5, y2, y2, y3, y4, y5 Coordinates
+        """
+
+        height, width, _ = frame.shape
+        stage_status = StageStatus(width=width, height=height)
+        m = 12 / self._min_face_size
+        min_layer = np.amin([height, width]) * m
+        scales = self.__compute_scale_pyramid(m, min_layer)
+
+        # We pipe here each of the stages
+        total_boxes, stage_status = self.__stage1(frame, scales, stage_status)
+        total_boxes, stage_status = self.__stage2(frame, total_boxes, stage_status)
+        bboxes, points = self.__stage3(frame, total_boxes, stage_status)
+
+        # Sort by location (to prevent flickering)
+        sort_idx = np.argsort(bboxes[:, 0])
+        bboxes = bboxes[sort_idx]
+        points = points[sort_idx]
+
+        # Transform to better shape and points now inside bbox
+        detections = []
+        cnt = 0
+        for i in range(bboxes.shape[0]):
+            conf = bboxes[i, -1].astype(np.float32)
+            if conf > self.min_detections_conf:
+                bboxes_c = np.reshape(bboxes[i, :-1], [2, 2]).astype(np.float32)
+                points_c = np.reshape(points[i], [2, 5]).transpose().astype(np.float32)
+                detections.append(
+                    Detection(
+                        idx=cnt,
+                        bbox=list(bboxes_c),
+                        landmarks=list(points_c),
+                        confidence=conf,
+                    )
+                )
+                cnt += 1
+        return frame, detections
+
+    def __compute_scale_pyramid(self, m, min_layer):
+        scales = []
+        factor_count = 0
+
+        while min_layer >= 12:
+            scales += [m * np.power(self._scale_factor, factor_count)]
+            min_layer = min_layer * self._scale_factor
+            factor_count += 1
+
+        return scales
+
+    @staticmethod
+    def __scale_image(image, scale: float):
+        """
+        Scales the image to a given scale.
+        :param image:
+        :param scale:
+        :return:
+        """
+        height, width, _ = image.shape
+
+        width_scaled = int(np.ceil(width * scale))
+        height_scaled = int(np.ceil(height * scale))
+
+        im_data = cv2.resize(image, (width_scaled, height_scaled), interpolation=cv2.INTER_AREA)
+
+        # Normalize the image's pixels
+        im_data_normalized = (im_data - 127.5) * 0.0078125
+
+        return im_data_normalized
+
+    @staticmethod
+    def __generate_bounding_box(imap, reg, scale, t):
+        # use heatmap to generate bounding boxes
+        stride = 2
+        cellsize = 12
+
+        imap = np.transpose(imap)
+        dx1 = np.transpose(reg[:, :, 0])
+        dy1 = np.transpose(reg[:, :, 1])
+        dx2 = np.transpose(reg[:, :, 2])
+        dy2 = np.transpose(reg[:, :, 3])
+
+        y, x = np.where(imap >= t)
+
+        if y.shape[0] == 1:
+            dx1 = np.flipud(dx1)
+            dy1 = np.flipud(dy1)
+            dx2 = np.flipud(dx2)
+            dy2 = np.flipud(dy2)
+
+        score = imap[(y, x)]
+        reg = np.transpose(np.vstack([dx1[(y, x)], dy1[(y, x)], dx2[(y, x)], dy2[(y, x)]]))
+
+        if reg.size == 0:
+            reg = np.empty(shape=(0, 3))
+
+        bb = np.transpose(np.vstack([y, x]))
+
+        q1 = np.fix((stride * bb + 1) / scale)
+        q2 = np.fix((stride * bb + cellsize) / scale)
+        boundingbox = np.hstack([q1, q2, np.expand_dims(score, 1), reg])
+
+        return boundingbox, reg
+
+    @staticmethod
+    def __nms(boxes, threshold, method):
+        """
+        Non Maximum Suppression.
+
+        :param boxes: np array with bounding boxes.
+        :param threshold:
+        :param method: NMS method to apply. Available values ('Min', 'Union')
+        :return:
+        """
+        if boxes.size == 0:
+            return np.empty((0, 3))
+
+        x1 = boxes[:, 0]
+        y1 = boxes[:, 1]
+        x2 = boxes[:, 2]
+        y2 = boxes[:, 3]
+        s = boxes[:, 4]
+
+        area = (x2 - x1 + 1) * (y2 - y1 + 1)
+        sorted_s = np.argsort(s)
+
+        pick = np.zeros_like(s, dtype=np.int16)
+        counter = 0
+        while sorted_s.size > 0:
+            i = sorted_s[-1]
+            pick[counter] = i
+            counter += 1
+            idx = sorted_s[0:-1]
+
+            xx1 = np.maximum(x1[i], x1[idx])
+            yy1 = np.maximum(y1[i], y1[idx])
+            xx2 = np.minimum(x2[i], x2[idx])
+            yy2 = np.minimum(y2[i], y2[idx])
+
+            w = np.maximum(0.0, xx2 - xx1 + 1)
+            h = np.maximum(0.0, yy2 - yy1 + 1)
+
+            inter = w * h
+
+            if method == "Min":
+                o = inter / np.minimum(area[i], area[idx])
+            else:
+                o = inter / (area[i] + area[idx] - inter)
+
+            sorted_s = sorted_s[np.where(o <= threshold)]
+
+        pick = pick[0:counter]
+
+        return pick
+
+    @staticmethod
+    def __pad(total_boxes, w, h):
+        # compute the padding coordinates (pad the bounding boxes to square)
+        tmp_w = (total_boxes[:, 2] - total_boxes[:, 0] + 1).astype(np.int32)
+        tmp_h = (total_boxes[:, 3] - total_boxes[:, 1] + 1).astype(np.int32)
+        numbox = total_boxes.shape[0]
+
+        dx = np.ones(numbox, dtype=np.int32)
+        dy = np.ones(numbox, dtype=np.int32)
+        edx = tmp_w.copy().astype(np.int32)
+        edy = tmp_h.copy().astype(np.int32)
+
+        x = total_boxes[:, 0].copy().astype(np.int32)
+        y = total_boxes[:, 1].copy().astype(np.int32)
+        ex = total_boxes[:, 2].copy().astype(np.int32)
+        ey = total_boxes[:, 3].copy().astype(np.int32)
+
+        tmp = np.where(ex > w)
+        edx.flat[tmp] = np.expand_dims(-ex[tmp] + w + tmp_w[tmp], 1)
+        ex[tmp] = w
+
+        tmp = np.where(ey > h)
+        edy.flat[tmp] = np.expand_dims(-ey[tmp] + h + tmp_h[tmp], 1)
+        ey[tmp] = h
+
+        tmp = np.where(x < 1)
+        dx.flat[tmp] = np.expand_dims(2 - x[tmp], 1)
+        x[tmp] = 1
+
+        tmp = np.where(y < 1)
+        dy.flat[tmp] = np.expand_dims(2 - y[tmp], 1)
+        y[tmp] = 1
+
+        return dy, edy, dx, edx, y, ey, x, ex, tmp_w, tmp_h
+
+    @staticmethod
+    def __rerec(bbox):
+        # convert bbox to square
+        height = bbox[:, 3] - bbox[:, 1]
+        width = bbox[:, 2] - bbox[:, 0]
+        max_side_length = np.maximum(width, height)
+        bbox[:, 0] = bbox[:, 0] + width * 0.5 - max_side_length * 0.5
+        bbox[:, 1] = bbox[:, 1] + height * 0.5 - max_side_length * 0.5
+        bbox[:, 2:4] = bbox[:, 0:2] + np.transpose(np.tile(max_side_length, (2, 1)))
+        return bbox
+
+    @staticmethod
+    def __bbreg(boundingbox, reg):
+        # calibrate bounding boxes
+        if reg.shape[1] == 1:
+            reg = np.reshape(reg, (reg.shape[2], reg.shape[3]))
+
+        w = boundingbox[:, 2] - boundingbox[:, 0] + 1
+        h = boundingbox[:, 3] - boundingbox[:, 1] + 1
+        b1 = boundingbox[:, 0] + reg[:, 0] * w
+        b2 = boundingbox[:, 1] + reg[:, 1] * h
+        b3 = boundingbox[:, 2] + reg[:, 2] * w
+        b4 = boundingbox[:, 3] + reg[:, 3] * h
+        boundingbox[:, 0:4] = np.transpose(np.vstack([b1, b2, b3, b4]))
+        return boundingbox
+
+    def __stage1(self, image, scales: list, stage_status: StageStatus):
+        """
+        First stage of the MTCNN.
+        :param image:
+        :param scales:
+        :param stage_status:
+        :return:
+        """
+        total_boxes = np.empty((0, 9))
+        status = stage_status
+
+        for scale in scales:
+            scaled_image = self.__scale_image(image, scale)
+
+            img_x = np.expand_dims(scaled_image, 0)
+            img_y = np.transpose(img_x, (0, 2, 1, 3))
+
+            out = tflite_inference(self.p_net, img_y)
+
+            out0 = np.transpose(out[0], (0, 2, 1, 3))
+            out1 = np.transpose(out[1], (0, 2, 1, 3))
+
+            boxes, _ = self.__generate_bounding_box(
+                out1[0, :, :, 1].copy(),
+                out0[0, :, :, :].copy(),
+                scale,
+                self._steps_threshold[0],
+            )
+
+            # inter-scale nms
+            pick = self.__nms(boxes.copy(), 0.5, "Union")
+            if boxes.size > 0 and pick.size > 0:
+                boxes = boxes[pick, :]
+                total_boxes = np.append(total_boxes, boxes, axis=0)
+
+        numboxes = total_boxes.shape[0]
+
+        if numboxes > 0:
+            pick = self.__nms(total_boxes.copy(), 0.7, "Union")
+            total_boxes = total_boxes[pick, :]
+
+            regw = total_boxes[:, 2] - total_boxes[:, 0]
+            regh = total_boxes[:, 3] - total_boxes[:, 1]
+
+            qq1 = total_boxes[:, 0] + total_boxes[:, 5] * regw
+            qq2 = total_boxes[:, 1] + total_boxes[:, 6] * regh
+            qq3 = total_boxes[:, 2] + total_boxes[:, 7] * regw
+            qq4 = total_boxes[:, 3] + total_boxes[:, 8] * regh
+
+            total_boxes = np.transpose(np.vstack([qq1, qq2, qq3, qq4, total_boxes[:, 4]]))
+            total_boxes = self.__rerec(total_boxes.copy())
+
+            total_boxes[:, 0:4] = np.fix(total_boxes[:, 0:4]).astype(np.int32)
+            status = StageStatus(
+                self.__pad(total_boxes.copy(), stage_status.width, stage_status.height),
+                width=stage_status.width,
+                height=stage_status.height,
+            )
+
+        return total_boxes, status
+
+    def __stage2(self, img, total_boxes, stage_status: StageStatus):
+        """
+        Second stage of the MTCNN.
+        :param img:
+        :param total_boxes:
+        :param stage_status:
+        :return:
+        """
+
+        num_boxes = total_boxes.shape[0]
+        if num_boxes == 0:
+            return total_boxes, stage_status
+
+        # second stage
+        tempimg = np.zeros(shape=(24, 24, 3, num_boxes))
+
+        for k in range(0, num_boxes):
+            tmp = np.zeros((int(stage_status.tmp_h[k]), int(stage_status.tmp_w[k]), 3))
+
+            tmp[
+                stage_status.dy[k] - 1 : stage_status.edy[k],
+                stage_status.dx[k] - 1 : stage_status.edx[k],
+                :,
+            ] = img[
+                stage_status.y[k] - 1 : stage_status.ey[k],
+                stage_status.x[k] - 1 : stage_status.ex[k],
+                :,
+            ]
+
+            if tmp.shape[0] > 0 and tmp.shape[1] > 0 or tmp.shape[0] == 0 and tmp.shape[1] == 0:
+                tempimg[:, :, :, k] = cv2.resize(tmp, (24, 24), interpolation=cv2.INTER_AREA)
+
+            else:
+                return np.empty(shape=(0,)), stage_status
+
+        tempimg = (tempimg - 127.5) * 0.0078125
+        tempimg1 = np.transpose(tempimg, (3, 1, 0, 2))
+
+        out = tflite_inference(self.r_net, tempimg1)
+
+        out0 = np.transpose(out[0])
+        out1 = np.transpose(out[1])
+
+        score = out1[1, :]
+
+        ipass = np.where(score > self._steps_threshold[1])
+
+        total_boxes = np.hstack([total_boxes[ipass[0], 0:4].copy(), np.expand_dims(score[ipass].copy(), 1)])
+
+        mv = out0[:, ipass[0]]
+
+        if total_boxes.shape[0] > 0:
+            pick = self.__nms(total_boxes, 0.7, "Union")
+            total_boxes = total_boxes[pick, :]
+            total_boxes = self.__bbreg(total_boxes.copy(), np.transpose(mv[:, pick]))
+            total_boxes = self.__rerec(total_boxes.copy())
+
+        return total_boxes, stage_status
+
+    def __stage3(self, img, total_boxes, stage_status: StageStatus):
+        """
+        Third stage of the MTCNN.
+
+        :param img:
+        :param total_boxes:
+        :param stage_status:
+        :return:
+        """
+        num_boxes = total_boxes.shape[0]
+        if num_boxes == 0:
+            return total_boxes, np.empty(shape=(0,))
+
+        total_boxes = np.fix(total_boxes).astype(np.int32)
+
+        status = StageStatus(
+            self.__pad(total_boxes.copy(), stage_status.width, stage_status.height),
+            width=stage_status.width,
+            height=stage_status.height,
+        )
+
+        tempimg = np.zeros((48, 48, 3, num_boxes))
+
+        for k in range(0, num_boxes):
+            tmp = np.zeros((int(status.tmp_h[k]), int(status.tmp_w[k]), 3))
+
+            tmp[status.dy[k] - 1 : status.edy[k], status.dx[k] - 1 : status.edx[k], :] = img[
+                status.y[k] - 1 : status.ey[k], status.x[k] - 1 : status.ex[k], :
+            ]
+
+            if tmp.shape[0] > 0 and tmp.shape[1] > 0 or tmp.shape[0] == 0 and tmp.shape[1] == 0:
+                tempimg[:, :, :, k] = cv2.resize(tmp, (48, 48), interpolation=cv2.INTER_AREA)
+            else:
+                return np.empty(shape=(0,)), np.empty(shape=(0,))
+
+        tempimg = (tempimg - 127.5) * 0.0078125
+        tempimg1 = np.transpose(tempimg, (3, 1, 0, 2))
+
+        out = tflite_inference(self.o_net, tempimg1)
+        out0 = np.transpose(out[0])
+        out1 = np.transpose(out[1])
+        out2 = np.transpose(out[2])
+
+        score = out2[1, :]
+
+        points = out1
+
+        ipass = np.where(score > self._steps_threshold[2])
+
+        points = points[:, ipass[0]]
+
+        total_boxes = np.hstack([total_boxes[ipass[0], 0:4].copy(), np.expand_dims(score[ipass].copy(), 1)])
+
+        mv = out0[:, ipass[0]]
+
+        w = total_boxes[:, 2] - total_boxes[:, 0] + 1
+        h = total_boxes[:, 3] - total_boxes[:, 1] + 1
+
+        points[0:5, :] = np.tile(w, (5, 1)) * points[0:5, :] + np.tile(total_boxes[:, 0], (5, 1)) - 1
+        points[5:10, :] = np.tile(h, (5, 1)) * points[5:10, :] + np.tile(total_boxes[:, 1], (5, 1)) - 1
+
+        if total_boxes.shape[0] > 0:
+            total_boxes = self.__bbreg(total_boxes.copy(), np.transpose(mv))
+            pick = self.__nms(total_boxes.copy(), 0.7, "Min")
+            total_boxes = total_boxes[pick, :]
+            points = points[:, pick]
+
+        return total_boxes, points.transpose()

tools/face_recognition.py

@@ -1,203 +1,114 @@
-from .nametypes import Detection, Identity
+from .utils import tflite_inference
+from .nametypes import Identity, Match
+from sklearn.metrics.pairwise import cosine_distances
 import numpy as np
 import cv2
-from sklearn.metrics.pairwise import cosine_distances
 from skimage.transform import SimilarityTransform
-
-
-def detect_faces(frame, model):
-    # Process the frame with MediaPipe Face Mesh
-    results = model.process(frame)
-
-    # Get the Bounding Boxes from the detected faces
-    detections = []
-    if results.multi_face_landmarks:
-        for face in results.multi_face_landmarks:
-            xs = [landmark.x for landmark in face.landmark]
-            ys = [landmark.y for landmark in face.landmark]
-            bbox = [min(xs), min(ys), max(xs), max(ys)]
-
-            FIVE_LANDMARKS = [470, 475, 1, 57, 287]
-
-            landmarks = [
-                [face.landmark[i].x, face.landmark[i].y] for i in FIVE_LANDMARKS
+from .utils import get_file
+import tflite_runtime.interpreter as tflite
+from typing import Literal
+
+
+BASE_URL = "https://github.com/Martlgap/FaceIDLight/releases/download/v.0.1/"
+
+FILE_HASHES = {
+    "mobileNet": "6c19b789f661caa8da735566490bfd8895beffb2a1ec97a56b126f0539991aa6",
+    "resNet": "f4d8b0194957a3ad766135505fc70a91343660151a8103bbb6c3b8ac34dbb4e2",
+}
+
+
+class FaceRecognition:
+    def __init__(
+        self,
+        min_similarity: float = 0.67,
+        model_name: Literal["mobileNet", "resNet50"] = "mobileNet",
+    ):
+        self.min_similarity = min_similarity
+        self.model = tflite.Interpreter(model_path=get_file(BASE_URL + f"{model_name}.tflite", FILE_HASHES[model_name]))
+
+    def __call__(self, frame, detections):
+        # Align Faces
+        faces, faces_aligned = [], []
+        for detection in detections:
+            face = frame[
+                int(detection.bbox[0][1]) : int(detection.bbox[1][1]),
+                int(detection.bbox[0][0]) : int(detection.bbox[1][0]),
             ]
+            try:
+                face = cv2.resize(face, (112, 112))
+            except:
+                face = np.zeros((112, 112, 3))
 
-            detections.append(Detection(bbox=bbox, landmarks=landmarks))
-    return detections
-
-
-def align(img, landmarks, target_size=(112, 112)):
-    # Transform to Landmark-Coordinates from relative landmark positions
-    dst = np.asarray(landmarks) * img.shape[:2][::-1]
-
-    # Target Landmarks-Coordinates from ArcFace Paper
-    src = np.array(
-        [
-            [38.2946, 51.6963],
-            [73.5318, 51.5014],
-            [56.0252, 71.7366],
-            [41.5493, 92.3655],
-            [70.7299, 92.2041],
-        ],
-        dtype=np.float32,
-    )
+            faces.append(face)
+            faces_aligned.append(self.align(frame, detection.landmarks))
 
-    # Estimate the transformation matrix
-    tform = SimilarityTransform()
-    tform.estimate(dst, src)
-    tmatrix = tform.params[0:2, :]
+        # Do Inference
+        if len(faces_aligned) == 0:
+            return []
 
-    # Apply the transformation matrix
-    img = cv2.warpAffine(img, tmatrix, target_size, borderValue=0.0)
+        # Normalize images from [0, 255] to [0, 1]
+        faces_aligned_norm = np.asarray(faces_aligned).astype(np.float32) / 255.0
 
-    return img
-
-
-def align_faces(img, detections):
-    updated_detections = []
-    for detection in detections:
-        updated_detections.append(
-            detection._replace(face=align(img, detection.landmarks))
-        )
-    return updated_detections
-
-
-# TODO Error when uploading image while running!
-def inference(detections, model):
-    updated_detections = []
-    faces = [detection.face for detection in detections if detection.face is not None]
-
-    if len(faces) > 0:
-        faces = np.asarray(faces).astype(np.float32) / 255
-        model.resize_tensor_input(model.get_input_details()[0]["index"], faces.shape)
-        model.allocate_tensors()
-        model.set_tensor(model.get_input_details()[0]["index"], faces)
-        model.invoke()
-        embs = [model.get_tensor(elem["index"]) for elem in model.get_output_details()][
-            0
-        ]
+        embs_det = tflite_inference(self.model, faces_aligned_norm)
+        embs_det = np.asarray(embs_det[0])
 
+        # Save Identities
+        identities = []
         for idx, detection in enumerate(detections):
-            updated_detections.append(detection._replace(embedding=embs[idx]))
-    return updated_detections
-
-
-def recognize_faces(detections, gallery, thresh=0.67):
-    if len(gallery) == 0 or len(detections) == 0:
-        return detections
-
-    gallery_embs = np.asarray([identity.embedding for identity in gallery])
-    detection_embs = np.asarray([detection.embedding for detection in detections])
-
-    cos_distances = cosine_distances(detection_embs, gallery_embs)
-
-    updated_detections = []
-    for idx, detection in enumerate(detections):
-        idx_min = np.argmin(cos_distances[idx])
-        if thresh and cos_distances[idx][idx_min] > thresh:
-            dist = cos_distances[idx][idx_min]
-            pred = None
-        else:
-            dist = cos_distances[idx][idx_min]
-            pred = idx_min
-        updated_detections.append(
-            detection._replace(
-                name=gallery[pred]
-                .name.split(".jpg")[0]
-                .split(".png")[0]
-                .split(".jpeg")[0]
-                if pred is not None
-                else None,
-                embedding_match=gallery[pred].embedding if pred is not None else None,
-                face_match=gallery[pred].image if pred is not None else None,
-                distance=dist,
-            )
-        )
-
-    return updated_detections
-
-
-def process_gallery(files, face_detection_model, face_recognition_model):
-    gallery = []
-    for file in files:
-        file_bytes = np.asarray(bytearray(file.read()), dtype=np.uint8)
-        img = cv2.cvtColor(
-            cv2.imdecode(file_bytes, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB
-        )
-
-        detections = detect_faces(img, face_detection_model)
-
-        # We accept only one face per image!
-        if detections == []:
-            continue
-        elif len(detections) > 1:
-            detections = detections[:1]
-
-        detections = align_faces(img, detections)
-        detections = inference(detections, face_recognition_model)
-
-        gallery.append(
-            Identity(
-                name=file.name,
-                embedding=detections[0].embedding,
-                image=detections[0].face,
-            )
-        )
-
-    return gallery
-
-
-def draw_detections(
-    frame, detections, bbox=True, landmarks=True, name=True):
-    shape = np.asarray(frame.shape[:2][::-1])
-
-    for detection in detections:
-        # Draw Landmarks
-        if landmarks:
-            for landmark in detection.landmarks:
-                cv2.circle(
-                    frame,
-                    (np.asarray(landmark) * shape).astype(int),
-                    2,
-                    (0, 0, 255),
-                    -1,
+            identities.append(
+                Identity(
+                    detection_idx=detection.idx,
+                    embedding=embs_det[idx],
+                    face_aligned=faces_aligned[idx],
                 )
-
-        # Draw Bounding Box
-        if bbox:
-            cv2.rectangle(
-                frame,
-                (np.asarray(detection.bbox[:2]) * shape).astype(int),
-                (np.asarray(detection.bbox[2:]) * shape).astype(int),
-                (0, 255, 0),
-                2,
-            )
-
-        # Draw Name
-        if name:
-            cv2.rectangle(
-                frame,
-                (
-                    int(detection.bbox[0] * shape[0]),
-                    int(detection.bbox[1] * shape[1] - (shape[1] // 25)),
-                ),
-                (int(detection.bbox[2] * shape[0]), int(detection.bbox[1] * shape[1])),
-                (255, 255, 255),
-                -1,
-            )
-
-            cv2.putText(
-                frame,
-                detection.name,
-                (
-                    int(detection.bbox[0] * shape[0] + shape[0] // 400),
-                    int(detection.bbox[1] * shape[1] - shape[1] // 100),
-                ),
-                cv2.LINE_AA,
-                0.5,
-                (0, 0, 0),
-                2,
             )
+        return identities
+
+    def find_matches(self, identities, gallery):
+        if len(gallery) == 0 or len(identities) == 0:
+            return []
+
+        # Get Embeddings
+        embs_gal = np.asarray([identity.embedding for identity in gallery])
+        embs_det = np.asarray([identity.embedding for identity in identities])
+
+        # Calculate Cosine Distances
+        cos_distances = cosine_distances(embs_det, embs_gal)
+
+        # Find Matches
+        matches = []
+        for ident_idx, identity in enumerate(identities):
+            dist_to_identity = cos_distances[ident_idx]
+            idx_min = np.argmin(dist_to_identity)
+            if dist_to_identity[idx_min] < self.min_similarity:
+                matches.append(
+                    Match(
+                        identity_idx=identity.detection_idx,
+                        gallery_idx=idx_min,
+                        distance=dist_to_identity[idx_min],
+                        name=gallery[idx_min].name,
+                    )
+                )
 
-    return frame
+        # Sort Matches by identity_idx
+        matches = sorted(matches, key=lambda match: match.gallery_idx)
+
+        return matches
+
+    @staticmethod
+    def align(img, landmarks_source, target_size=(112, 112)):
+        landmarks_target = np.array(
+            [
+                [38.2946, 51.6963],
+                [73.5318, 51.5014],
+                [56.0252, 71.7366],
+                [41.5493, 92.3655],
+                [70.7299, 92.2041],
+            ],
+            dtype=np.float32,
+        )
+        tform = SimilarityTransform()
+        tform.estimate(landmarks_source, landmarks_target)
+        tmatrix = tform.params[0:2, :]
+        face_aligned = cv2.warpAffine(img, tmatrix, target_size, borderValue=0.0)
+        return face_aligned

tools/nametypes.py

@@ -3,14 +3,17 @@ import numpy as np
 
 
 class Detection(NamedTuple):
-    bbox: List[int]
-    landmarks: List[List[int]]
+    idx: int = None
+    bbox: List[List[float]] = None
+    landmarks: List[List[float]] = None
+    confidence: float = None
+
+
+class Identity(NamedTuple):
+    detection_idx: int = None
     name: str = None
-    face: np.ndarray = None
     embedding: np.ndarray = None
-    embedding_match: np.ndarray = None
-    face_match: np.ndarray = None
-    distance: float = None
+    face_aligned: np.ndarray = None
 
 
 class Stats(NamedTuple):
@@ -18,13 +21,13 @@ class Stats(NamedTuple):
     resolution: List[int] = [None, None, None]
     num_faces: int = 0
     detection: float = None
-    alignment: float = None
-    inference: float = None
     recognition: float = None
-    drawing: float = None
+    matching: float = None
+    annotation: float = None
 
 
-class Identity(NamedTuple):
-    name: str
-    embedding: np.ndarray
-    image: np.ndarray
+class Match(NamedTuple):
+    identity_idx: int = None
+    gallery_idx: int = None
+    distance: float = None
+    name: str = None

tools/pca.py

@@ -0,0 +1,59 @@
+from sklearn.decomposition import PCA
+import numpy as np
+import plotly.express as px
+
+
+def pca(matches, identities, gallery, dim=3):
+    """
+    Perform PCA on embeddings.
+    Args:
+        embeddings: np.array of shape (n_embeddings, 512)
+    Returns:
+        embeddings_pca: np.array of shape (n_embeddings, 3)
+    """
+
+    # Get Gallery and Detection Embeddings and stich them together in groups
+    embeddings = np.concatenate(
+        [[gallery[match.gallery_idx].embedding, identities[match.identity_idx].embedding] for match in matches],
+        axis=0,
+    )
+
+    # Get Identity Names and stich them together in groups
+    identity_names = np.concatenate(
+        [[gallery[match.gallery_idx].name, gallery[match.gallery_idx].name] for match in matches],
+        axis=0,
+    )
+
+    # Do 3D PCA
+    pca = PCA(n_components=dim)
+    pca.fit(embeddings)
+    embeddings_pca = pca.transform(embeddings)
+
+    if dim == 3:
+        fig = px.scatter_3d(
+            embeddings_pca,
+            x=0,
+            y=1,
+            z=2,
+            opacity=0.7,
+            color=identity_names,
+            color_discrete_sequence=px.colors.qualitative.Vivid,
+        )
+        fig.update_traces(marker=dict(size=4))
+    elif dim == 2:
+        fig = px.scatter(
+            embeddings_pca,
+            x=0,
+            y=1,
+            opacity=0.7,
+            color=identity_names,
+            color_discrete_sequence=px.colors.qualitative.Vivid,
+        )
+        fig.update_traces(marker=dict(size=4))
+        fig.update_xaxes(showgrid=True)
+        fig.update_yaxes(showgrid=True)
+    else:
+        raise ValueError("dim must be either 2 or 3")
+    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))
+
+    return fig

tools/utils.py

@@ -1,13 +1,16 @@
 import logging
 import os
-import urllib.request
-from pathlib import Path
 import streamlit as st
 from twilio.rest import Client
 import os
-import cv2
 import numpy as np
 import hashlib
+import tempfile
+import os
+import hashlib
+from tqdm import tqdm
+from zipfile import ZipFile
+from urllib.request import urlopen
 
 
 logger = logging.getLogger(__name__)
@@ -25,9 +28,7 @@ def get_ice_servers(name="twilio"):
             account_sid = os.environ["TWILIO_ACCOUNT_SID"]
             auth_token = os.environ["TWILIO_AUTH_TOKEN"]
         except KeyError:
-            logger.warning(
-                "Twilio credentials are not set. Fallback to a free STUN server from Google."
-            )
+            logger.warning("Twilio credentials are not set. Fallback to a free STUN server from Google.")
             return [{"urls": ["stun:stun.l.google.com:19302"]}]
 
         client = Client(account_sid, auth_token)
@@ -41,9 +42,7 @@ def get_ice_servers(name="twilio"):
             username = os.environ["METERED_USERNAME"]
             credential = os.environ["METERED_CREDENTIAL"]
         except KeyError:
-            logger.warning(
-                "Metered credentials are not set. Fallback to a free STUN server from Google."
-            )
+            logger.warning("Metered credentials are not set. Fallback to a free STUN server from Google.")
             return [{"urls": ["stun:stun.l.google.com:19302"]}]
 
         ice_servers = [
@@ -78,70 +77,6 @@ def get_ice_servers(name="twilio"):
         raise ValueError(f"Unknown name: {name}")
 
 
-def get_hash(filepath):
-    hasher = hashlib.sha256()
-    with open(filepath, "rb") as file:
-        for chunk in iter(lambda: file.read(65535), b""):
-            hasher.update(chunk)
-    return hasher.hexdigest()
-
-
-def download_file(url, model_path: Path, file_hash=None):
-    if model_path.exists():
-        if file_hash:
-            hasher = hashlib.sha256()
-            with open(model_path, "rb") as file:
-                for chunk in iter(lambda: file.read(65535), b""):
-                    hasher.update(chunk)
-            if not hasher.hexdigest() == file_hash:
-                print(
-                    "A local file was found, but it seems to be incomplete or outdated because the file hash does not "
-                    "match the original value of "
-                    + file_hash
-                    + " so data will be downloaded."
-                )
-                download = True
-            else:
-                print("Using a verified local file.")
-                download = False
-    else:
-        model_path.mkdir(parents=True, exist_ok=True)
-        print("Downloading data ...")
-        download = True
-
-    if download:
-        # These are handles to two visual elements to animate.
-        weights_warning, progress_bar = None, None
-        try:
-            weights_warning = st.warning("Downloading %s..." % url)
-            progress_bar = st.progress(0)
-            with open(model_path, "wb") as output_file:
-                with urllib.request.urlopen(url) as response:
-                    length = int(response.info()["Content-Length"])
-                    counter = 0.0
-                    MEGABYTES = 2.0**20.0
-                    while True:
-                        data = response.read(8192)
-                        if not data:
-                            break
-                        counter += len(data)
-                        output_file.write(data)
-
-                        # We perform animation by overwriting the elements.
-                        weights_warning.warning(
-                            "Downloading %s... (%6.2f/%6.2f MB)"
-                            % (url, counter / MEGABYTES, length / MEGABYTES)
-                        )
-                        progress_bar.progress(min(counter / length, 1.0))
-
-        # Finally, we remove these visual elements by calling .empty().
-        finally:
-            if weights_warning is not None:
-                weights_warning.empty()
-            if progress_bar is not None:
-                progress_bar.empty()
-
-
 # Function to format floats within a list
 def format_dflist(val):
     if isinstance(val, list):
@@ -156,20 +91,74 @@ def format_dflist(val):
         return val
 
 
-def display_match(d):
-    im = np.concatenate([d.face, d.face_match])
-    border_size = 2
-    border = cv2.copyMakeBorder(
-        im,
-        top=border_size,
-        bottom=border_size,
-        left=border_size,
-        right=border_size,
-        borderType=cv2.BORDER_CONSTANT,
-        value=(255, 255, 120),
-    )
-    return border
-
-
 def rgb(r, g, b):
     return "#{:02x}{:02x}{:02x}".format(r, g, b)
+
+
+def tflite_inference(model, img):
+    """Inferences an image through the model with tflite interpreter on CPU
+    :param model: a tflite.Interpreter loaded with a model
+    :param img: image
+    :return: list of outputs of the model
+    """
+    # Check if img is np.ndarray
+    if not isinstance(img, np.ndarray):
+        img = np.asarray(img)
+
+    # Check if dim is 4
+    if len(img.shape) == 3:
+        img = np.expand_dims(img, axis=0)
+
+    input_details = model.get_input_details()
+    output_details = model.get_output_details()
+    model.resize_tensor_input(input_details[0]["index"], img.shape)
+    model.allocate_tensors()
+    model.set_tensor(input_details[0]["index"], img.astype(np.float32))
+    model.invoke()
+    return [model.get_tensor(elem["index"]) for elem in output_details]
+
+
+def get_file(origin, file_hash, is_zip=False):
+    tmp_file = os.path.join(tempfile.gettempdir(), "FaceIDLight", origin.split("/")[-1])
+    os.makedirs(os.path.dirname(tmp_file), exist_ok=True)
+    if not os.path.exists(tmp_file):
+        download = True
+    else:
+        hasher = hashlib.sha256()
+        with open(tmp_file, "rb") as file:
+            for chunk in iter(lambda: file.read(65535), b""):
+                hasher.update(chunk)
+        if not hasher.hexdigest() == file_hash:
+            print(
+                "A local file was found, but it seems to be incomplete or outdated because the file hash does not "
+                "match the original value of " + file_hash + " so data will be downloaded."
+            )
+            download = True
+        else:
+            download = False
+
+    if download:
+        response = urlopen(origin)
+        with tqdm.wrapattr(
+            open(tmp_file, "wb"),
+            "write",
+            miniters=1,
+            desc="Downloading " + origin.split("/")[-1] + " to: " + tmp_file,
+            total=getattr(response, "length", None),
+        ) as file:
+            for chunk in response:
+                file.write(chunk)
+            file.close()
+    if is_zip:
+        with ZipFile(tmp_file, "r") as zipObj:
+            zipObj.extractall(tmp_file.split(".")[0])
+        tmp_file = os.path.join(tmp_file.split(".")[0])
+    return tmp_file
+
+
+def get_hash(filepath):
+    hasher = hashlib.sha256()
+    with open(filepath, "rb") as file:
+        for chunk in iter(lambda: file.read(65535), b""):
+            hasher.update(chunk)
+    return hasher.hexdigest()