Spaces:
Running
Running
| import os | |
| import cv2 | |
| from torch.utils.model_zoo import load_url | |
| from ..core import FaceDetector | |
| from .net_s3fd import s3fd | |
| from .bbox import * | |
| from .detect import * | |
| models_urls = { | |
| 's3fd': 'https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth', | |
| } | |
| class SFDDetector(FaceDetector): | |
| def __init__(self, device, path_to_detector=os.path.join(os.path.dirname(os.path.abspath(__file__)), 's3fd.pth'), verbose=False): | |
| super(SFDDetector, self).__init__(device, verbose) | |
| # Initialise the face detector | |
| if not os.path.isfile(path_to_detector): | |
| model_weights = load_url(models_urls['s3fd']) | |
| else: | |
| model_weights = torch.load(path_to_detector) | |
| self.face_detector = s3fd() | |
| self.face_detector.load_state_dict(model_weights) | |
| self.face_detector.to(device) | |
| self.face_detector.eval() | |
| 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) | |
| 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, images): | |
| bboxlists = batch_detect(self.face_detector, images, device=self.device) | |
| keeps = [nms(bboxlists[:, i, :], 0.3) for i in range(bboxlists.shape[1])] | |
| bboxlists = [bboxlists[keep, i, :] for i, keep in enumerate(keeps)] | |
| bboxlists = [[x for x in bboxlist if x[-1] > 0.5] for bboxlist in bboxlists] | |
| return bboxlists | |
| def reference_scale(self): | |
| return 195 | |
| def reference_x_shift(self): | |
| return 0 | |
| def reference_y_shift(self): | |
| return 0 | |