github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-08-22 06:23:20 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

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

Browse source 
Updated S3FDExtractor
This commit is contained in:
Jeremy Hummel 2019-08-23 13:51:55 -07:00 • committed by GitHub
commit 953cc81db3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 37 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

61
facelib/S3FDExtractor.py
View file

@ -3,26 +3,39 @@ from pathlib import Path
import cv2 import cv2
from nnlib import nnlib from nnlib import nnlib
class S3FDExtractor(object): class S3FDExtractor(object):
"""
S3FD: Single Shot Scale-invariant Face Detector
https://arxiv.org/pdf/1708.05237.pdf
"""
def __init__(self): def __init__(self):
exec( nnlib.import_all(), locals(), globals() ) exec(nnlib.import_all(), locals(), globals())
model_path = Path(__file__).parent / "S3FD.h5" model_path = Path(__file__).parent / "S3FD.h5"
if not model_path.exists(): 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): def __enter__(self):
return self return self
def __exit__(self, exc_type=None, exc_value=None, traceback=None): 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: if is_bgr:
input_image = input_image[:,:,::-1] input_image = input_image[:, :, ::-1]
is_bgr = False is_bgr = False
(h, w, ch) = input_image.shape (h, w, ch) = input_image.shape
@ -32,35 +45,36 @@ class S3FDExtractor(object):
scale_to = max(64, scale_to) scale_to = max(64, scale_to)
input_scale = d / 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 = [] detected_faces = []
for ltrb in self.refine (olist): for ltrb in self._refine(olist, nms_thresh):
l,t,r,b = [ x*input_scale for x in ltrb] l, t, r, b = [x * input_scale for x in ltrb]
bt = b-t bt = b - t
if min(r-l,bt) < 40: #filtering faces < 40pix by any side if min(r - l, bt) < 40: # filtering faces < 40pix by any side
continue continue
b += bt*0.1 #enlarging bottom line a bit for 2DFAN-4, because default is not enough covering a chin 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) ] ) detected_faces.append([int(x) for x in (l, t, r, b)])
return detected_faces return detected_faces
def refine(self, olist): def _refine(self, olist, thresh):
bboxlist = [] bboxlist = []
for i, ((ocls,), (oreg,)) in enumerate ( zip ( olist[::2], olist[1::2] ) ): for i, ((ocls,), (oreg,)) in enumerate(zip(olist[::2], olist[1::2])):
stride = 2**(i + 2) # 4,8,16,32,64,128 stride = 2 ** (i + 2) # 4,8,16,32,64,128
s_d2 = stride / 2 s_d2 = stride / 2
s_m4 = stride * 4 s_m4 = stride * 4
for hindex, windex in zip(*np.where(ocls > 0.05)): for hindex, windex in zip(*np.where(ocls > 0.05)):
score = ocls[hindex, windex] 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 = np.array([windex * stride + s_d2, hindex * stride + s_d2, s_m4, s_m4])
priors_2p = priors[2:] priors_2p = priors[2:]
box = np.concatenate((priors[:2] + loc[:2] * 0.1 * priors_2p, 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:] / 2
box[2:] += box[:2] box[2:] += box[:2]
@ -69,12 +83,11 @@ class S3FDExtractor(object):
bboxlist = np.array(bboxlist) bboxlist = np.array(bboxlist)
if len(bboxlist) == 0: if len(bboxlist) == 0:
bboxlist = np.zeros((1, 5)) bboxlist = np.zeros((1, 5))
#Originally 0.3 thresh bboxlist = bboxlist[self._refine_nms(bboxlist, 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 = [ x[:-1].astype(np.int) for x in bboxlist if x[-1] >= 0.5]
return bboxlist return bboxlist
def refine_nms(self, dets, thresh): def _refine_nms(self, dets, nms_thresh):
keep = list() keep = list()
if len(dets) == 0: if len(dets) == 0:
return keep 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) 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) 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] order = order[inds + 1]
return keep return keep