Replace Face Detection Library

datasets/vggface2.py

@@ -23,8 +23,8 @@
 from torchvision import transforms
 
 import cv2
-import face_detection
 import kornia.geometry.transform as GT
+from batch_face import RetinaFace
 from PIL import Image
 from skimage import transform as trans
 from tqdm import tqdm
@@ -99,8 +99,11 @@ def __extract_gt(self):
         """
         Extracts the ground truth from the dataset
         """
-        detector = face_detection.build_detector("RetinaNetResNet50", confidence_threshold=.5,
-                                                 nms_iou_threshold=.4)
+        if torch.cuda.is_available():
+            detector = RetinaFace(gpu_id=torch.cuda.current_device(), network="resnet50")
+        else:
+            detector = RetinaFace(gpu_id=-1, network="resnet50")
+
         img_paths = list(glob.glob(os.path.join(self.d_path + '/**/', '*.jpg'), recursive=True))
         nf_number = 0
         words_count = 0
@@ -111,22 +114,17 @@ def __extract_gt(self):
             boxes = []
             image = cv2.imread(jpg)
 
-            img_max = max(image.shape[0], image.shape[1])
-            if img_max > 1320:
-                continue
-            bboxes, lndmrks = detector.batched_detect_with_landmarks(np.expand_dims(image, 0))
-            bboxes = bboxes[0]
-            lndmrks = lndmrks[0]
+            faces = detector(image)
 
-            if (bboxes.shape[0] == 0) or (lndmrks.shape[0] == 0):
+            if len(faces) == 0:
                 nf_number += 1
                 continue
 
-            for box in bboxes:
+            for face in faces:
+                box = face[0]
                 box = np.clip(box[:4], 0, None)
                 boxes.append(box)
-
-            lndmrks = lndmrks[0]
+            lndmrks = faces[0][1]
 
             dir_name = os.path.dirname(jpg)
             lbl = os.path.relpath(dir_name, self.d_path)

requirements.txt

@@ -6,7 +6,7 @@ Pillow>=7
 PyYAML>=5.1.1
 albumentations>=1.3.0
 faiss-cpu==1.7.4
-face-detection==0.2.2
+batch-face>=1.4.0
 h5py>=3.7.0
 kornia==0.6.8
 librosa>=0.7.2