Implement batch scoring and enhance existing methods

face_image_quality.py

@@ -120,6 +120,188 @@ def apply_mtcnn(self, face_image : np.ndarray):
 
 
 
+    def get_score(self, aligned_img : np.ndarray, 
+                        T : int = 100,
+                        alpha : float = 130.0,
+                        r : float = 0.88):
+        """
+        Calculates the SER-FIQ score for a given aligned image using T passes.
+
+
+        Parameters
+        ----------
+        aligned_img : np.ndarray, shape (3, h, w)
+            Aligned face image, in RGB format.
+        T : int, optional
+            Amount of forward passes to use. The default is 100.
+        alpha : float, optional
+            Stretching factor, can be choosen to scale the score values
+        r : float, optional
+            Score displacement
+
+        Returns
+        -------
+        SER-FIQ score : float.
+
+        """
+        # Color Channel is not the first dimension, swap dims.
+        if aligned_img.shape[0] != 3:
+            aligned_img = np.transpose(aligned_img, (2,0,1))
+
+        input_blob = np.expand_dims(aligned_img, axis=0)
+        repeated = np.repeat(input_blob, T, axis=0)
+        gpu_repeated = mx.nd.array(repeated, ctx=self.device)
+
+        X = self.insightface(gpu_repeated).asnumpy()
+
+        norm = normalize(X, axis=1)
+
+        # Only get the upper triangle of the distance matrix
+        eucl_dist = euclidean_distances(norm, norm)[np.triu_indices(T, k=1)]
+
+        # Calculate score as given in the paper
+        score = 2*(1/(1+np.exp(np.mean(eucl_dist))))
+        # Normalize value based on alpha and r
+        return 1 / (1+np.exp(-(alpha * (score - r))))
+
+    def get_scores_batch(self, aligned_imgs: list, 
+                              T: int = 100, 
+                              alpha: float = 130.0, 
+                              r: float = 0.88):
+        """
+        Calculates SER-FIQ scores for a list of aligned images efficiently by 
+        stacking their forward passes together into a single network invocation.
+
+        Parameters
+        ----------
+        aligned_imgs : list of np.ndarray
+            A list containing aligned face images.
+        T : int, optional
+            Amount of forward passes to use per image.
+        alpha : float, optional
+            Stretching factor.
+        r : float, optional
+            Score displacement.
+
+        Returns
+        -------
+        scores : list of float
+            List of calculated SER-FIQ scores matching the input sequence.
+        """
+        if not aligned_imgs:
+            return []
+
+        processed_imgs = [
+            np.transpose(img, (2, 0, 1)) if img.shape[0] != 3 else img 
+            for img in aligned_imgs
+        ]
+
+        # Shape: (Num_Images * T, 3, H, W)
+        batch_blob = np.concatenate([np.repeat(np.expand_dims(img, axis=0), T, axis=0) for img in processed_imgs], axis=0)
+        gpu_batch = mx.nd.array(batch_blob, ctx=self.device)
+
+        all_features = self.insightface(gpu_batch).asnumpy()
+        normalized_features = normalize(all_features, axis=1)
+
+        scores = []
+        triu_idx = np.triu_indices(T, k=1)
+
+        for i in range(len(aligned_imgs)):
+            # Extract features belonging to the i-th image block
+            img_features = normalized_features[i * T : (i + 1) * T]
+            eucl_dist = euclidean_distances(img_features, img_features)[triu_idx]
+
+            score = 2 * (1 / (1 + np.exp(np.mean(eucl_dist))))
+            final_score = 1 / (1 + np.exp(-(alpha * (score - r))))
+            scores.append(final_score)
+
+        return scores
+                 gpu:int=0, # Which gpu should be used -> gpu id
+                 det:int=0, # Mtcnn option, 1= Use R+O, 0=Detect from beginning
+                 ):
+        """
+        Reimplementing Insightface's FaceModel class.
+        Now the dropout output and the network output are returned after a forward pass.
+
+        Parameters
+        ----------
+        gpu : int, optional
+            The GPU to be used by Mxnet. The default is 0.
+            If set to None, CPU is used instead.
+        det : int, optional
+            Mtcnn option, 1= Use R+0, 0= Detect from beginning. The default is 0.
+
+        Returns
+        -------
+        None.
+
+        """
+
+        if gpu is None:
+            self.device = mx.cpu()
+        else:
+            self.device = mx.gpu(gpu)
+
+        self.insightface = gluon.nn.SymbolBlock.imports(
+                                    "./insightface/model/insightface-symbol.json",
+                                    ['data'],
+                                    "./insightface/model/insightface-0000.params", 
+                                    ctx=self.device
+                           )
+
+
+        self.det_minsize = 50
+        self.det_threshold = [0.6,0.7,0.8]
+        self.det = det
+
+        self.preprocess = face_preprocess.preprocess
+
+        thrs = self.det_threshold if det==0 else [0.0,0.0,0.2]
+
+        self.detector = mtcnn_detector.MtcnnDetector(model_folder="./insightface/mtcnn-model/", 
+                                                    ctx=self.device, 
+                                                    num_worker=1, 
+                                                    accurate_landmark = True, 
+                                                    threshold=thrs
+                                                    )
+
+    def apply_mtcnn(self, face_image : np.ndarray):
+        """
+        Applies MTCNN Detector on the given face image and returns
+        the cropped image.
+
+        If no face could be detected None is returned.
+
+        Parameters
+        ----------
+        face_image : np.ndarray
+            Face imaged loaded via OpenCV.
+
+        Returns
+        -------
+        Face Image : np.ndarray, shape (3,112,112).
+        None, if no face could be detected
+
+        """
+        detected = self.detector.detect_face(face_image, det_type=self.det)
+
+        if detected is None:
+            return None
+
+        bbox, points = detected
+
+        if bbox.shape[0] == 0:
+            return None
+
+        points = points[0, :].reshape((2,5)).T
+
+        image = self.preprocess(face_image, bbox, points, image_size="112,112")
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        return np.transpose(image, (2,0,1))
+
+
+
     def get_score(self, aligned_img : np.ndarray, 
                         T : int = 100,
                         alpha : float = 130.0,