|
from .core import FaceDetector
|
|
from .detect import *
|
|
|
|
|
|
class SFDDetector(FaceDetector):
|
|
def __init__(self, path_to_detector=None, device="cuda", verbose=False):
|
|
super(SFDDetector, self).__init__(device, verbose)
|
|
self.device = device
|
|
self.face_detector = s3fd()
|
|
self.face_detector.load_state_dict(torch.load(path_to_detector))
|
|
self.face_detector.eval()
|
|
if self.device == "cuda":
|
|
self.face_detector.cuda()
|
|
|
|
def detect_from_image(self, tensor_or_path):
|
|
image = self.tensor_or_path_to_ndarray(tensor_or_path)
|
|
bboxlist = detect(self.face_detector, image, device=self.device)[0]
|
|
keep = nms(bboxlist, 0.3)
|
|
bboxlist = bboxlist[keep, :]
|
|
bboxlist = [x for x in bboxlist if x[-1] > 0.5]
|
|
|
|
return bboxlist
|
|
|
|
def detect_from_batch(self, tensor):
|
|
bboxlists = batch_detect(self.face_detector, tensor, device=self.device)
|
|
error = False
|
|
new_bboxlists = []
|
|
error_index = -1
|
|
for i in range(bboxlists.shape[0]):
|
|
bboxlist = bboxlists[i]
|
|
keep = nms(bboxlist, 0.3)
|
|
if len(keep) > 0:
|
|
bboxlist = bboxlist[keep, :]
|
|
bboxlist = [x for x in bboxlist if x[-1] > 0.5]
|
|
new_bboxlists.append(bboxlist)
|
|
else:
|
|
error = True
|
|
error_index = i
|
|
new_bboxlists.append([])
|
|
|
|
return new_bboxlists, error, error_index
|
|
|
|
@property
|
|
def reference_scale(self):
|
|
return 195
|
|
|
|
@property
|
|
def reference_x_shift(self):
|
|
return 0
|
|
|
|
@property
|
|
def reference_y_shift(self):
|
|
return 0
|
|
|
