diff --git a/face_image_quality.py b/face_image_quality.py index c3d36c5..65064f4 100644 --- a/face_image_quality.py +++ b/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,