add retinaface

.gitignore

@@ -10,9 +10,7 @@ __pycache__
 *.npz
 *.hdf5
 
-src/face_models/model-ir-v1
-src/face_models/model-mtcnn
-src/face_models/model-r100-ii
+src/face_models/models
 
 images/rs
 images/rs_aligned

requirements.txt

@@ -19,6 +19,7 @@ idna==2.6
 imageio==2.6.1
 imbalanced-learn==0.6.1
 imblearn==0.0
+insightface===0.1.5
 isort==4.3.21
 joblib==0.14.1
 Keras-Applications==1.0.8
@@ -28,7 +29,7 @@ lazy-object-proxy==1.4.3
 Markdown==3.1.1
 matplotlib==3.1.2
 mccabe==0.6.1
-mxnet==1.5.0
+mxnet==1.6.0
 networkx==2.4
 numpy==1.16.4
 oauthlib==3.1.0
@@ -46,7 +47,7 @@ pyparsing==2.4.6
 python-dateutil==2.8.1
 pytz==2019.3
 PyWavelets==1.1.1
-requests==2.20.0
+requests==2.18.4
 requests-oauthlib==1.3.0
 rsa==4.0
 scikit-image==0.16.2

src/biocapsule.py

@@ -1,38 +1,56 @@
-import os
 import numpy as np
 from scipy.signal import convolve2d
-from argparse import ArgumentParser
 
 
 class BioCapsuleGenerator:
     def __signature_extraction(self, feature):
-        lvl1 = convolve2d(feature.reshape(32, 16), np.ones(
-            (5, 5)) / 25., mode="same", boundary="wrap")
+        lvl1 = convolve2d(
+            feature.reshape(32, 16),
+            np.ones((5, 5)) / 25.0,
+            mode="same",
+            boundary="wrap",
+        )
+
         lvl2 = feature.reshape(32, 16) - lvl1
-        signature = np.around(np.average(lvl2, axis=1) * 100.).astype(int) % 9
+
+        signature = np.around(np.average(lvl2, axis=1) * 100.0).astype(int) % 9
+
         return signature
 
     def __key_generation(self, signature):
         key = np.empty((0,))
+
         for sig in signature:
             np.random.seed(sig)
             key = np.append(key, np.random.choice(2, 16))
+
         key = (key * 2) - 1
+
         return key
 
     def biocapsule(self, user_feature, rs_feature):
         user_signature = self.__signature_extraction(user_feature)
         user_key = self.__key_generation(user_signature)
+
         rs_signature = self.__signature_extraction(rs_feature)
         rs_key = self.__key_generation(rs_signature)
-        return np.multiply(user_feature, rs_key) + np.multiply(rs_feature, user_key)
+
+        bc = np.multiply(user_feature, rs_key) + np.multiply(
+            rs_feature, user_key
+        )
+
+        return bc
 
     def biocapsule_batch(self, user_features, rs_feature):
         rs_signature = self.__signature_extraction(rs_feature)
         rs_key = self.__key_generation(rs_signature)
+
         for i in range(user_features.shape[0]):
             user_signature = self.__signature_extraction(user_features[i])
             user_key = self.__key_generation(user_signature)
-            user_features[i] = np.multiply(user_features[i], rs_key) + \
-                np.multiply(rs_feature, user_key)
+
+            user_features[i] = np.multiply(
+                user_features[i], rs_key
+            ) + np.multiply(rs_feature, user_key)
+
         return user_features

src/face.py

@@ -18,13 +18,13 @@ class FaceNet:
 
     """
 
-    def __init__(self, gpu: int = -1):
+    def __init__(self, gpu: int = -1, detector: str = "mtcnn"):
         """Initialize FaceNet model object. Provide an int
         corresponding to a GPU id to use for models. If -1 is given
         CPU is used rather than GPU.
 
         """
-        self.__model = FaceNetModel(gpu)
+        self.__model = FaceNetModel(gpu, detector)
 
     def preprocess(self, face_img: np.ndarray) -> np.ndarray:
         """Preprocess facial image for FaceNet feature extraction
@@ -66,15 +66,18 @@ def extract(self, face_img: np.ndarray, align: bool = True) -> np.ndarray:
             face_img = self.preprocess(face_img)
 
         if face_img.shape != (160, 160, 3):
+            if len(face_img.shape) == 2:
+                face_img = cv2.cvtColor(face_img, cv2.COLOR_GRAY2BGR)
+
             face_img = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)
             face_img = cv2.resize(face_img, (160, 160))
 
         return self.__model.get_feature(face_img)
 
 
 class ArcFace:
-    def __init__(self, gpu: int = -1):
-        self.__model = ArcFaceModel(gpu)
+    def __init__(self, gpu: int = -1, detector: str = "mtcnn"):
+        self.__model = ArcFaceModel(gpu, detector)
 
     def preprocess(self, face_img: np.ndarray) -> np.ndarray:
         """Preprocess facial image for ArcFace feature extraction
@@ -115,6 +118,9 @@ def extract(self, face_img: np.ndarray, align: bool = True) -> np.ndarray:
             face_img = self.preprocess(face_img)
 
         if face_img.shape != (3, 112, 112):
+            if len(face_img.shape) == 2:
+                face_img = cv2.cvtColor(face_img, cv2.COLOR_GRAY2BGR)
+
             face_img = cv2.resize(face_img, (112, 112))
             face_img = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)
             face_img = np.rollaxis(face_img, 2, 0)
@@ -123,34 +129,36 @@ def extract(self, face_img: np.ndarray, align: bool = True) -> np.ndarray:
 
 
 def extract_dataset(
-    dataset: str, extractor: str = "arcface", gpu: int = -1
-) -> np.ndarray:
+    dataset: str,
+    method: str = "arcface",
+    detector: str = "mtcnn",
+    flipped: bool = True,
+    gpu: int = -1,
+):
     """Extract feature vectors of each image within a dataset.
-    Return array conatining all extracted features.
+    Save array conatining all extracted features to disk.
 
     Parameters
     ----------
     dataset: str
         Dataset to extract features from. Examples would be gtdb or lfw
-    extractor: str = "arcface"
+    method: str = "arcface"
         Model to use for feature extraction. Currently supported options are
         arcface/facenet
+    detector: str = "mtcnn"
+        Model to use for facial preprocessing. Currently supported options are
+        mtcnn/retinaface
+    flipped: bool = True
+        Flag denoting if flipped features should be extracted
     gpu: int = -1
         GPU id to use for feature extraction and preprocessing models. If -1
         is given, CPU is used rather than GPU
 
-    Returns
-    -------
-    np.ndarray
-        Array of features corresponding each image from a dataset. Subject ids
-        are appended to end of feature vectors. Resulting output will be of
-        shape (number of dataset images)x513
-
     """
-    if extractor == "arcface":
-        face = ArcFace(gpu)
+    if method == "arcface":
+        face = ArcFace(gpu, detector)
     else:
-        face = FaceNet(gpu)
+        face = FaceNet(gpu, detector)
 
     dataset_path = f"images/{dataset}"
 
@@ -161,27 +169,53 @@ def extract_dataset(
         os.listdir(dataset_path), key=lambda subject: subject.lower()
     )
 
-    image_cnt = 0
+    img_cnt = 0
     for subject_id, subject in enumerate(subjects):
-        progress_bar(f"{dataset} {extractor}", (image_cnt + 1) / file_cnt)
+        progress_bar(f"{dataset} {method}", (img_cnt + 1) / file_cnt)
 
         for image in os.listdir(f"{dataset_path}/{subject}"):
-            image = cv2.imread(f"{dataset_path}/{subject}/{image}")
+            img = cv2.imread(f"{dataset_path}/{subject}/{image}")
+
+            feature = face.extract(img)
+            features[img_cnt, :] = np.append(feature, subject_id + 1)
+
+            img_cnt += 1
+
+    np.savez_compressed(
+        f"data/{dataset}_{method}_{detector}_feat.npz", features
+    )
+
+    if flipped:
+        flipped_features = np.zeros((file_cnt, 513))
+
+        img_cnt = 0
+        for subject_id, subject in enumerate(subjects):
+            progress_bar(
+                f"{dataset} {method} flipped", (img_cnt + 1) / file_cnt
+            )
 
-            feature = face.extract(image)
-            features[image_cnt, :] = np.append(feature, subject_id + 1)
+            for image in os.listdir(f"{dataset_path}/{subject}"):
+                img = cv2.imread(f"{dataset_path}/{subject}/{image}")
+                img = cv2.flip(img, 1)
 
-            image_cnt += 1
+                flipped_feature = face.extract(img)
+                flipped_features[img_cnt, :] = np.append(
+                    flipped_feature, subject_id + 1
+                )
 
-    return features
+                img_cnt += 1
+
+        np.savez_compressed(
+            f"data/{dataset}_{method}_{detector}_flip_feat.npz", features
+        )
 
 
 if __name__ == "__main__":
     """
-    facenet = FaceNet()
+    facenet = FaceNet(-1, "retinaface")
 
-    img_1 = cv2.imread("src/face_models/examples/tom1.jpg"))
-    img_2 = cv2.imread("src/face_models/examples/adrien.jpg"))
+    img_1 = cv2.imread("src/face_models/examples/tom1.jpg")
+    img_2 = cv2.imread("src/face_models/examples/adrien.jpg")
 
     feat_1 = facenet.extract(img_1)
     feat_2 = facenet.extract(img_2)
@@ -194,7 +228,7 @@ def extract_dataset(
     cv2.waitKey(0)
     cv2.destroyAllWindows()
 
-    arcface = ArcFace()
+    arcface = ArcFace(-1, "mtcnn")
 
     feat_1 = arcface.extract(img_1)
     feat_2 = arcface.extract(img_2)
@@ -204,13 +238,11 @@ def extract_dataset(
     cv2.imshow("before", img_2)
     img = arcface.preprocess(img_2)
     cv2.imshow(
-        "after",
-        cv2.cvtColor(np.rollaxis(img, 0, 3), cv2.COLOR_RGB2BGR)
+        "after", cv2.cvtColor(np.rollaxis(img, 0, 3), cv2.COLOR_RGB2BGR)
     )
     cv2.waitKey(0)
     cv2.destroyAllWindows()
     """
-
     parser = ArgumentParser()
     parser.add_argument(
         "-d",
@@ -225,6 +257,21 @@ def extract_dataset(
         choices=["arcface", "facenet"],
         help="method to use in feature extraction",
     )
+    parser.add_argument(
+        "-det",
+        "--detector",
+        required=True,
+        choices=["mtcnn", "retinaface"],
+        help="method to use in facial preprocessing",
+    )
+    parser.add_argument(
+        "-f",
+        "--flipped",
+        required=False,
+        action="store_true",
+        default=False,
+        help="extract features for flipped versions of images",
+    )
     parser.add_argument(
         "-gpu",
         "--gpu",
@@ -235,11 +282,10 @@ def extract_dataset(
     )
     args = vars(parser.parse_args())
 
-    features = extract_dataset(args["dataset"], args["method"], args["gpu"])
-
-    np.savez_compressed(
-        os.path.join(
-            "data/{}_{}_feat.npz".format(args["dataset"], args["method"])
-        ),
-        features,
+    features = extract_dataset(
+        args["dataset"],
+        args["method"],
+        args["detector"],
+        args["flipped"],
+        args["gpu"],
     )

src/face_models/__init__.py

@@ -1 +1 @@
-from face_models.face_model import ArcFaceModel, FaceNetModel
+from .face_model import ArcFaceModel, FaceNetModel