378 lines
19 KiB
Python
378 lines
19 KiB
Python
"""
|
|
@File : test.py
|
|
@Author : Songkey
|
|
@Email : songkey@pku.edu.cn
|
|
@Date : 11/1/2024
|
|
@Desc : Created by Shengjie Wu (wu.shengjie@immomo.com)
|
|
"""
|
|
|
|
import numpy as np
|
|
from .utils import OneEuroFilter
|
|
|
|
def cult_dis(old_kpts, new_kpts):
|
|
dis = np.sqrt(np.square(new_kpts[:, 0] - old_kpts[:, 0]) + np.square(new_kpts[:, 1] - old_kpts[:, 1]))
|
|
return dis
|
|
|
|
class Smoother222(object):
|
|
def __init__(self):
|
|
# face config
|
|
self.face_idx = list(range(0, 33))
|
|
self.face_down_idx = list(range(9, 24))
|
|
self.filter_face = OneEuroFilter()
|
|
# nose config
|
|
self.nose_idx = list(range(33, 48))
|
|
self.filter_nose = OneEuroFilter()
|
|
# eyebrow config
|
|
self.eyebrow_idx = list(range(48, 74))
|
|
self.filter_eyebrow = OneEuroFilter()
|
|
# eye config
|
|
self.left_eye_idx = list(range(74, 96))
|
|
self.filter_left_eye = OneEuroFilter()
|
|
self.right_eye_idx = list(range(96, 118))
|
|
self.filter_right_eye = OneEuroFilter()
|
|
# mouth config
|
|
self.mouth_idx = list(range(118, 182))
|
|
self.filter_mouth = OneEuroFilter()
|
|
# pupil config
|
|
self.left_pupil_idx = list(range(182, 202))
|
|
self.filter_left_pupil = OneEuroFilter()
|
|
self.right_pupil_idx = list(range(202, 222))
|
|
self.filter_right_pupil = OneEuroFilter()
|
|
self.prev_points = None
|
|
|
|
def smooth(self, new_points, face_dis):
|
|
if self.prev_points is None:
|
|
self.prev_points = new_points.copy()
|
|
return new_points
|
|
dis = cult_dis(self.prev_points, new_points) / face_dis
|
|
smooth_points = new_points.copy()
|
|
|
|
# smooth face
|
|
if np.mean(dis[self.face_down_idx]) < 0.005:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.face_down_idx]) / 0.005
|
|
fcmin_tmp = 0.05 * ratio_tmp
|
|
beta_tmp = 0.05 * ratio_tmp
|
|
smooth_points[self.face_idx] = self.filter_face(new_points[self.face_idx], self.prev_points[self.face_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.face_down_idx]) < 0.02:
|
|
# filter
|
|
ratio_tmp = (np.mean(dis[self.face_down_idx]) - 0.005) / (0.02 - 0.005)
|
|
fcmin_tmp = 0.05 + (0.3 - 0.05) * ratio_tmp
|
|
beta_tmp = 0.05 + (0.3 - 0.05) * ratio_tmp
|
|
smooth_points[self.face_idx] = self.filter_face(new_points[self.face_idx], self.prev_points[self.face_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# filter
|
|
smooth_points[self.face_idx] = self.filter_face(new_points[self.face_idx], self.prev_points[self.face_idx],
|
|
fcmin=0.3, beta=0.3)
|
|
# smooth nose
|
|
if np.mean(dis[self.nose_idx]) < 0.003:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.nose_idx]) / 0.003
|
|
fcmin_tmp = 0.03 * ratio_tmp
|
|
beta_tmp = 0.03 * ratio_tmp
|
|
smooth_points[self.nose_idx] = self.filter_nose(new_points[self.nose_idx], self.prev_points[self.nose_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.nose_idx]) < 0.02:
|
|
ratio_tmp = (np.mean(dis[self.nose_idx]) - 0.003) / (0.02 - 0.003)
|
|
fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
smooth_points[self.nose_idx] = self.filter_nose(new_points[self.nose_idx], self.prev_points[self.nose_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# filter
|
|
smooth_points[self.nose_idx] = self.filter_nose(new_points[self.nose_idx], self.prev_points[self.nose_idx],
|
|
fcmin=0.7, beta=0.7)
|
|
# smooth eyebrow
|
|
# print(np.mean(dis[self.eyebrow_idx]))
|
|
if np.mean(dis[self.eyebrow_idx]) < 0.003:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.eyebrow_idx]) / 0.003
|
|
fcmin_tmp = 0.02 * ratio_tmp
|
|
beta_tmp = 0.02 * ratio_tmp
|
|
smooth_points[self.eyebrow_idx] = self.filter_eyebrow(new_points[self.eyebrow_idx],
|
|
self.prev_points[self.eyebrow_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.eyebrow_idx]) < 0.02:
|
|
# filter
|
|
ratio_tmp = (np.mean(dis[self.eyebrow_idx]) - 0.003) / (0.02 - 0.003)
|
|
fcmin_tmp = 0.02 + (0.5 - 0.02) * ratio_tmp
|
|
beta_tmp = 0.02 + (0.5 - 0.02) * ratio_tmp
|
|
smooth_points[self.eyebrow_idx] = self.filter_eyebrow(new_points[self.eyebrow_idx],
|
|
self.prev_points[self.eyebrow_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# filter
|
|
smooth_points[self.eyebrow_idx] = self.filter_eyebrow(new_points[self.eyebrow_idx],
|
|
self.prev_points[self.eyebrow_idx],
|
|
fcmin=0.5, beta=0.5)
|
|
# smooth eye
|
|
if np.mean(dis[self.left_eye_idx]) < 0.003:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.left_eye_idx]) / 0.003
|
|
fcmin_tmp = 0.03 * ratio_tmp
|
|
beta_tmp = 0.03 * ratio_tmp
|
|
smooth_points[self.left_eye_idx] = self.filter_left_eye(new_points[self.left_eye_idx],
|
|
self.prev_points[self.left_eye_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.left_eye_idx]) < 0.02:
|
|
# filter
|
|
ratio_tmp = (np.mean(dis[self.left_eye_idx]) - 0.003) / (0.02 - 0.003)
|
|
fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
smooth_points[self.left_eye_idx] = self.filter_left_eye(new_points[self.left_eye_idx],
|
|
self.prev_points[self.left_eye_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# fast
|
|
smooth_points[self.left_eye_idx] = self.filter_left_eye(new_points[self.left_eye_idx],
|
|
self.prev_points[self.left_eye_idx],
|
|
fcmin=0.7, beta=0.7)
|
|
if np.mean(dis[self.right_eye_idx]) < 0.003:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.right_eye_idx]) / 0.003
|
|
fcmin_tmp = 0.03 * ratio_tmp
|
|
beta_tmp = 0.03 * ratio_tmp
|
|
smooth_points[self.right_eye_idx] = self.filter_right_eye(new_points[self.right_eye_idx],
|
|
self.prev_points[self.right_eye_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.right_eye_idx]) < 0.02:
|
|
# filter
|
|
ratio_tmp = (np.mean(dis[self.right_eye_idx]) - 0.003) / (0.02 - 0.003)
|
|
fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
smooth_points[self.right_eye_idx] = self.filter_right_eye(new_points[self.right_eye_idx],
|
|
self.prev_points[self.right_eye_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# fast
|
|
smooth_points[self.right_eye_idx] = self.filter_right_eye(new_points[self.right_eye_idx],
|
|
self.prev_points[self.right_eye_idx],
|
|
fcmin=0.7, beta=0.7)
|
|
|
|
# smooth mouth
|
|
if np.mean(dis[self.mouth_idx]) < 0.003:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.mouth_idx]) / 0.003
|
|
fcmin_tmp = 0.05 * ratio_tmp
|
|
beta_tmp = 0.05 * ratio_tmp
|
|
smooth_points[self.mouth_idx] = self.filter_mouth(new_points[self.mouth_idx],
|
|
self.prev_points[self.mouth_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.mouth_idx]) < 0.02:
|
|
# filter
|
|
ratio_tmp = (np.mean(dis[self.mouth_idx]) - 0.003) / (0.02 - 0.003)
|
|
fcmin_tmp = 0.05 + (0.7 - 0.05) * ratio_tmp
|
|
beta_tmp = 0.05 + (0.7 - 0.05) * ratio_tmp
|
|
smooth_points[self.mouth_idx] = self.filter_mouth(new_points[self.mouth_idx],
|
|
self.prev_points[self.mouth_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# fast
|
|
smooth_points[self.mouth_idx] = self.filter_mouth(new_points[self.mouth_idx],
|
|
self.prev_points[self.mouth_idx],
|
|
fcmin=0.7, beta=0.7)
|
|
|
|
# smooth pupil
|
|
if np.mean(dis[self.left_pupil_idx]) < 0.003:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.left_pupil_idx]) / 0.003
|
|
fcmin_tmp = 0.03 * ratio_tmp
|
|
beta_tmp = 0.03 * ratio_tmp
|
|
smooth_points[self.left_pupil_idx] = self.filter_left_pupil(new_points[self.left_pupil_idx],
|
|
self.prev_points[self.left_pupil_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.left_pupil_idx]) < 0.02:
|
|
# filter
|
|
ratio_tmp = (np.mean(dis[self.left_pupil_idx]) - 0.003) / (0.02 - 0.003)
|
|
fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
smooth_points[self.left_pupil_idx] = self.filter_left_pupil(new_points[self.left_pupil_idx],
|
|
self.prev_points[self.left_pupil_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# fast
|
|
smooth_points[self.left_pupil_idx] = self.filter_left_pupil(new_points[self.left_pupil_idx],
|
|
self.prev_points[self.left_pupil_idx],
|
|
fcmin=0.7, beta=0.7)
|
|
if np.mean(dis[self.right_pupil_idx]) < 0.003:
|
|
# stable
|
|
ratio_tmp = np.mean(dis[self.right_pupil_idx]) / 0.003
|
|
fcmin_tmp = 0.03 * ratio_tmp
|
|
beta_tmp = 0.03 * ratio_tmp
|
|
smooth_points[self.right_pupil_idx] = self.filter_right_pupil(new_points[self.right_pupil_idx],
|
|
self.prev_points[self.right_pupil_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
elif np.mean(dis[self.right_pupil_idx]) < 0.02:
|
|
# filter
|
|
ratio_tmp = (np.mean(dis[self.right_pupil_idx]) - 0.003) / (0.02 - 0.003)
|
|
fcmin_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
beta_tmp = 0.03 + (0.7 - 0.03) * ratio_tmp
|
|
smooth_points[self.right_pupil_idx] = self.filter_right_pupil(new_points[self.right_pupil_idx],
|
|
self.prev_points[self.right_pupil_idx],
|
|
fcmin=fcmin_tmp, beta=beta_tmp)
|
|
else:
|
|
# fast
|
|
smooth_points[self.right_pupil_idx] = self.filter_right_pupil(new_points[self.right_pupil_idx],
|
|
self.prev_points[self.right_pupil_idx],
|
|
fcmin=0.7, beta=0.7)
|
|
|
|
# update pre points
|
|
self.prev_points = smooth_points
|
|
return smooth_points
|
|
|
|
class HelloCameraDemo(object):
|
|
def __init__(self, face_alignment_module, reset=False):
|
|
self.face_alignment_module = face_alignment_module
|
|
self.face_prob_th = 0.0001
|
|
self.min_face = 96
|
|
self.face_image_size = self.face_alignment_module.face_image_size
|
|
self.trackingFaces = []
|
|
self.reset = reset
|
|
|
|
def reset_track(self):
|
|
self.trackingFaces = []
|
|
|
|
def forward(self, src_image, reset=False, pre_rect=None):
|
|
if self.reset or reset:
|
|
self.trackingFaces = []
|
|
|
|
if len(self.trackingFaces) == 0:
|
|
if pre_rect is not None:
|
|
detected_faces = [pre_rect]
|
|
else:
|
|
detected_faces, _, _ = self.face_alignment_module.face_detector.detect(src_image)
|
|
for face_rect in detected_faces:
|
|
new_tracking_object = {
|
|
'face_rect': face_rect,
|
|
'rotate_angle': 0.0,
|
|
'pre_kpt_222': None,
|
|
'face_dis': np.sqrt(
|
|
np.square((face_rect[2] - face_rect[0])) + np.square((face_rect[3] - face_rect[1]))),
|
|
'smoother_222': Smoother222(),
|
|
'prob': 0
|
|
}
|
|
self.trackingFaces.append(new_tracking_object)
|
|
else:
|
|
detected_faces, _, _ = self.face_alignment_module.face_detector.detect(src_image)
|
|
for face_rect in detected_faces:
|
|
new_tracking_object = {
|
|
'face_rect': face_rect,
|
|
'rotate_angle': 0.0,
|
|
'pre_kpt_222': None,
|
|
'face_dis': np.sqrt(
|
|
np.square((face_rect[2] - face_rect[0])) + np.square((face_rect[3] - face_rect[1]))),
|
|
'smoother_222': Smoother222(),
|
|
'prob': 0
|
|
}
|
|
self.trackingFaces.append(new_tracking_object)
|
|
|
|
delete_idx_list = []
|
|
for face_idx, tracking_face in enumerate(self.trackingFaces):
|
|
if tracking_face['pre_kpt_222'] is not None:
|
|
result_dict = self.face_alignment_module.forward(src_image, pre_pts=tracking_face['pre_kpt_222'],
|
|
iterations=3)
|
|
else:
|
|
result_dict = self.face_alignment_module.forward(src_image, face_box=tracking_face['face_rect'],
|
|
iterations=3)
|
|
|
|
if result_dict['prob'] < self.face_prob_th:
|
|
if not face_idx in delete_idx_list:
|
|
delete_idx_list.append(face_idx)
|
|
continue
|
|
|
|
landmarks_final = tracking_face['smoother_222'].smooth(result_dict['pt222'], tracking_face['face_dis'])
|
|
tracking_face['pre_kpt_222'] = landmarks_final
|
|
|
|
left_eye_corner = landmarks_final[74]
|
|
right_eye_corner = landmarks_final[96]
|
|
|
|
radian = np.arctan2(right_eye_corner[1] - left_eye_corner[1],
|
|
right_eye_corner[0] - left_eye_corner[0] + 0.00000001)
|
|
rotate_angle = np.rad2deg(radian)
|
|
face_x_min, face_x_max = np.min(landmarks_final[:, 0]), np.max(landmarks_final[:, 0])
|
|
face_y_min, face_y_max = np.min(landmarks_final[:, 1]), np.max(landmarks_final[:, 1])
|
|
face_bbox = [face_x_min, face_y_min, face_x_max, face_y_max]
|
|
face_dis = np.linalg.norm(landmarks_final[0] - landmarks_final[32])
|
|
|
|
if face_x_max - face_x_min < self.min_face or face_y_max - face_y_min < self.min_face:
|
|
if not face_idx in delete_idx_list:
|
|
delete_idx_list.append(face_idx)
|
|
|
|
euler_pred = result_dict["euler_rad"]
|
|
pitch = np.rad2deg(euler_pred[0])
|
|
yaw = np.rad2deg(euler_pred[1])
|
|
roll = np.rad2deg(euler_pred[2])
|
|
# print("pitch, yaw, roll", pitch, yaw, roll)
|
|
|
|
# one filter model
|
|
max_euler = abs(pitch) + (abs(yaw)*0.6)
|
|
face_dis *= (1.0 + max_euler / 18.0)
|
|
|
|
# two filter model
|
|
tracking_face['face_rect'] = face_bbox
|
|
tracking_face['rotate_angle'] = rotate_angle
|
|
tracking_face['face_dis'] = face_dis
|
|
tracking_face['prob'] = result_dict['prob']
|
|
tracking_face['pitch'] = pitch
|
|
tracking_face['yaw'] = yaw
|
|
tracking_face['roll'] = roll
|
|
tracking_face['euler_rad'] = result_dict["euler_rad"]
|
|
|
|
if len(self.trackingFaces) > 1:
|
|
for face_idx, tracking_face_target in enumerate(self.trackingFaces):
|
|
if face_idx in delete_idx_list:
|
|
continue
|
|
for idx, tracking_face in enumerate(self.trackingFaces):
|
|
if idx in delete_idx_list:
|
|
continue
|
|
if face_idx == idx: continue
|
|
iou_temp = self.count_iou(tracking_face_target['face_rect'], tracking_face['face_rect'])
|
|
# prog 2
|
|
if iou_temp > 0.12:
|
|
if self.area(tracking_face_target['face_rect']) - self.area(tracking_face['face_rect']) < 0:
|
|
if not face_idx in delete_idx_list:
|
|
delete_idx_list.append(face_idx)
|
|
else:
|
|
if not idx in delete_idx_list:
|
|
delete_idx_list.append(idx)
|
|
|
|
|
|
idx_offset = 0
|
|
for delete_idx in sorted(delete_idx_list):
|
|
self.trackingFaces.pop(delete_idx - idx_offset)
|
|
idx_offset += 1
|
|
|
|
return self.trackingFaces
|
|
|
|
def count_iou(self, boxA, boxB):
|
|
# determine the (x, y)-coordinates of the intersection rectangle
|
|
xA = max(boxA[0], boxB[0])
|
|
yA = max(boxA[1], boxB[1])
|
|
xB = min(boxA[2], boxB[2])
|
|
yB = min(boxA[3], boxB[3])
|
|
|
|
# compute the area of intersection rectangle
|
|
interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))
|
|
if interArea == 0:
|
|
return 0
|
|
# compute the area of both the prediction and ground-truth
|
|
# rectangles
|
|
boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
|
|
boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))
|
|
|
|
# compute the intersection over union by taking the intersection
|
|
# area and dividing it by the sum of prediction + ground-truth
|
|
# areas - the interesection area
|
|
iou = interArea / float(boxAArea + boxBArea - interArea)
|
|
|
|
# return the intersection over union value
|
|
return iou
|
|
|
|
def area(self, bbox):
|
|
w = bbox[3]-bbox[1]
|
|
h = bbox[2]-bbox[0]
|
|
return w * h
|
|
|