File size: 2,516 Bytes
03da825 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
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)
|