diff --git a/facelib/S3FDExtractor.py b/facelib/S3FDExtractor.py new file mode 100644 index 0000000..3763cf2 --- /dev/null +++ b/facelib/S3FDExtractor.py @@ -0,0 +1,95 @@ +import numpy as np +from pathlib import Path +import cv2 +from nnlib import nnlib + +class S3FDExtractor(object): + def __init__(self): + self.scale_factor = 3 + + exec( nnlib.import_all(), locals(), globals() ) + + model_path = Path(__file__).parent / "S3FD.h5" + if not model_path.exists(): + return None + + 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 + + def extract_from_bgr (self, input_image): + input_image = input_image[:,:,::-1].copy() + (h, w, ch) = input_image.shape + + d = w if w > h else h + input_scale = d / ( (d // self.scale_factor) - (d % self.scale_factor) ) + 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) ) + + detected_faces = self.refine (olist) + + #enlarging bottom line a bit for 2DFAN-4, because default is not enough covering a chin + for face in detected_faces: + l,t,r,b = face + tb = (b-t) * 0.1 + face[3] += tb + + return [ (int(face[0]*input_scale), + int(face[1]*input_scale), + int(face[2]*input_scale), + int(face[3]*input_scale)) for face in detected_faces ] + + def refine(self, olist): + bboxlist = [] + 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, :] + 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)) ) + box[:2] -= box[2:] / 2 + box[2:] += box[:2] + + bboxlist.append([*box, score]) + + bboxlist = np.array(bboxlist) + if len(bboxlist) == 0: + bboxlist = np.zeros((1, 5)) + + bboxlist = bboxlist[self.refine_nms(bboxlist, 0.3), :] + bboxlist = [ x[:-1].astype(np.int) for x in bboxlist if x[-1] >= 0.5] + return bboxlist + + def refine_nms(self, dets, thresh): + keep = list() + if len(dets) == 0: + return keep + + x_1, y_1, x_2, y_2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4] + areas = (x_2 - x_1 + 1) * (y_2 - y_1 + 1) + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx_1, yy_1 = np.maximum(x_1[i], x_1[order[1:]]), np.maximum(y_1[i], y_1[order[1:]]) + xx_2, yy_2 = np.minimum(x_2[i], x_2[order[1:]]), np.minimum(y_2[i], y_2[order[1:]]) + + 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] + order = order[inds + 1] + return keep \ No newline at end of file