| import time | |
| import face_alignment | |
| import numpy as np | |
| import cv2 | |
| from DDFA_V2.FaceBoxes.FaceBoxes_ONNX import FaceBoxes_ONNX | |
| class FaceBoxEstimator: | |
| def __init__(self, device): | |
| device_idx = int(device.split(':')[-1]) | |
| providers = [ | |
| ('CUDAExecutionProvider', { | |
| 'device_id': device_idx, | |
| }), | |
| 'CPUExecutionProvider', | |
| ] | |
| self.face_boxes = FaceBoxes_ONNX(providers=providers) | |
| self.bboxes = [] | |
| def add_box(self, box): | |
| self.bboxes.append(box) | |
| if len(self.bboxes) > 5: | |
| self.bboxes.pop(0) | |
| def clear_history(self): | |
| self.bboxes = [] | |
| def get_current_box(self): | |
| return np.mean(self.bboxes, axis=0) | |
| def __call__(self, img): | |
| boxes = self.face_boxes(img) | |
| if len(boxes) == 0: | |
| return boxes | |
| self.add_box(boxes[0]) | |
| return [self.get_current_box()] | |
| class FaceDetector: | |
| def __init__(self, device): | |
| self.facebox_detector = FaceBoxEstimator(device) | |
| self.fa = face_alignment.FaceAlignment( | |
| face_alignment.LandmarksType.TWO_D, flip_input=False, device=device | |
| ) | |
| self.keypoints = [] | |
| def add(self, kps): | |
| self.keypoints.append(kps) | |
| if len(self.keypoints) > 5: | |
| self.keypoints.pop(0) | |
| def clear_history(self): | |
| self.bboxes = [] | |
| def get_current_keypoints(self): | |
| return np.mean(self.keypoints, axis=0) | |
| def __call__(self, img): | |
| img = np.array(img) | |
| scale_factor = 512 / max(img.shape[0], img.shape[1]) | |
| img = cv2.resize(img, (0, 0), fx=scale_factor, fy=scale_factor) | |
| bbox = self.facebox_detector(img) | |
| landmarks = self.fa.get_landmarks(img, detected_faces=bbox) | |
| if landmarks is None: | |
| return None, None | |
| landmarks = np.array(landmarks)[0] / scale_factor | |
| bbox = np.array(bbox)[0] / scale_factor | |
| self.add(landmarks) | |
| landmarks = self.get_current_keypoints() | |
| return landmarks, bbox | |
| if __name__ == '__main__': | |
| from PIL import Image | |
| detector = FaceDetector() | |
| img = Image.open('00000/00008.png') | |
| lm, bbox = detector(img) | |
| bbox = bbox.astype(int) | |
| print('Score:', bbox[-1]) | |
| img = np.array(img) | |
| cv2.rectangle(img, bbox[:2], bbox[2:4], (255, 0, 0), thickness=5) | |
| img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) | |
| img = cv2.resize(img, (0, 0), fx=0.2, fy=0.2) | |
| cv2.imshow('debug', img) | |
| cv2.waitKey(-1) | |