Merge pull request #31 from faceshiftlabs/feat/consistent-dpi/s3fd

Updated S3FDExtractor

facelib/S3FDExtractor.py

@@ -3,26 +3,39 @@ from pathlib import Path
 import cv2
 from nnlib import nnlib
 
+
 class S3FDExtractor(object):
+    """
+    S3FD: Single Shot Scale-invariant Face Detector
+    https://arxiv.org/pdf/1708.05237.pdf
+    """
     def __init__(self):
-        exec( nnlib.import_all(), locals(), globals() )
+        exec(nnlib.import_all(), locals(), globals())
 
         model_path = Path(__file__).parent / "S3FD.h5"
         if not model_path.exists():
-            return None
+            raise Exception(f'Could not find S3DF model at path {model_path}')
 
-        self.model = nnlib.keras.models.load_model ( str(model_path) )
+        self.model = nnlib.keras.models.load_model(str(model_path))
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
-        return False #pass exception between __enter__ and __exit__ to outter level
+        return False  # pass exception between __enter__ and __exit__ to outter level
 
-    def extract (self, input_image, is_bgr=True):
+    def extract(self, input_image, is_bgr=True, nms_thresh=0.3):
+        """
+        Extracts the bounding boxes for all faces found in image
+        :param input_image: The image to look for faces in
+        :param is_bgr: Is this image in OpenCV's BGR color mode, if not, assume RGB color mode
+        :param nms_thresh: The NMS (non-maximum suppression) threshold. Of all bounding boxes found, only return
+        bounding boxes with an overlap ratio less then threshold
+        :return:
+        """
 
         if is_bgr:
-            input_image = input_image[:,:,::-1]
+            input_image = input_image[:, :, ::-1]
             is_bgr = False
 
         (h, w, ch) = input_image.shape
@@ -32,35 +45,36 @@ class S3FDExtractor(object):
         scale_to = max(64, scale_to)
 
         input_scale = d / scale_to
-        input_image = cv2.resize (input_image, ( int(w/input_scale), int(h/input_scale) ), interpolation=cv2.INTER_LINEAR)
+        input_image = cv2.resize(input_image, (int(w / input_scale), int(h / input_scale)),
+                                 interpolation=cv2.INTER_LINEAR)
 
-        olist = self.model.predict( np.expand_dims(input_image,0) )
+        olist = self.model.predict(np.expand_dims(input_image, 0))
 
         detected_faces = []
-        for ltrb in self.refine (olist):
-            l,t,r,b = [ x*input_scale for x in ltrb]
-            bt = b-t
-            if min(r-l,bt) < 40: #filtering faces < 40pix by any side
+        for ltrb in self._refine(olist, nms_thresh):
+            l, t, r, b = [x * input_scale for x in ltrb]
+            bt = b - t
+            if min(r - l, bt) < 40:  # filtering faces < 40pix by any side
                 continue
-            b += bt*0.1 #enlarging bottom line a bit for 2DFAN-4, because default is not enough covering a chin
-            detected_faces.append ( [int(x) for x in (l,t,r,b) ] )
+            b += bt * 0.1  # enlarging bottom line a bit for 2DFAN-4, because default is not enough covering a chin
+            detected_faces.append([int(x) for x in (l, t, r, b)])
 
         return detected_faces
 
-    def refine(self, olist):
+    def _refine(self, olist, thresh):
         bboxlist = []
-        for i, ((ocls,), (oreg,)) in enumerate ( zip ( olist[::2], olist[1::2] ) ):
-            stride = 2**(i + 2)    # 4,8,16,32,64,128
+        for i, ((ocls,), (oreg,)) in enumerate(zip(olist[::2], olist[1::2])):
+            stride = 2 ** (i + 2)  # 4,8,16,32,64,128
             s_d2 = stride / 2
             s_m4 = stride * 4
 
             for hindex, windex in zip(*np.where(ocls > 0.05)):
                 score = ocls[hindex, windex]
-                loc   = oreg[hindex, windex, :]
+                loc = oreg[hindex, windex, :]
                 priors = np.array([windex * stride + s_d2, hindex * stride + s_d2, s_m4, s_m4])
                 priors_2p = priors[2:]
                 box = np.concatenate((priors[:2] + loc[:2] * 0.1 * priors_2p,
-                                      priors_2p * np.exp(loc[2:] * 0.2)) )
+                                      priors_2p * np.exp(loc[2:] * 0.2)))
                 box[:2] -= box[2:] / 2
                 box[2:] += box[:2]
 
@@ -69,12 +83,11 @@ class S3FDExtractor(object):
         bboxlist = np.array(bboxlist)
         if len(bboxlist) == 0:
             bboxlist = np.zeros((1, 5))
-        #Originally 0.3 thresh
-        bboxlist = bboxlist[self.refine_nms(bboxlist, 0.8), :]
-        bboxlist = [ x[:-1].astype(np.int) for x in bboxlist if x[-1] >= 0.5]
+        bboxlist = bboxlist[self._refine_nms(bboxlist, thresh), :]
+        bboxlist = [x[:-1].astype(np.int) for x in bboxlist if x[-1] >= 0.5]
         return bboxlist
 
-    def refine_nms(self, dets, thresh):
+    def _refine_nms(self, dets, nms_thresh):
         keep = list()
         if len(dets) == 0:
             return keep
@@ -93,6 +106,6 @@ class S3FDExtractor(object):
             width, height = np.maximum(0.0, xx_2 - xx_1 + 1), np.maximum(0.0, yy_2 - yy_1 + 1)
             ovr = width * height / (areas[i] + areas[order[1:]] - width * height)
 
-            inds = np.where(ovr <= thresh)[0]
+            inds = np.where(ovr <= nms_thresh)[0]
             order = order[inds + 1]
         return keep