From b400adadf5e78e6837526789cc60085be9fefc87 Mon Sep 17 00:00:00 2001 From: jh Date: Fri, 23 Aug 2019 13:41:12 -0700 Subject: [PATCH] Adds doc comments, explains the threshold parameter, cleaned up and refactored the class --- facelib/S3FDExtractor.py | 61 ++++++++++++++++++++++++---------------- 1 file changed, 37 insertions(+), 24 deletions(-) diff --git a/facelib/S3FDExtractor.py b/facelib/S3FDExtractor.py index 601521d..1372720 100644 --- a/facelib/S3FDExtractor.py +++ b/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