94 lines
4.3 KiB
Python
94 lines
4.3 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 cv2
|
|
import numpy as np
|
|
from .hello_face_det import HelloFaceDet
|
|
from .utils import get_warp_mat_bbox, get_warp_mat_bbox_by_gt_pts_float, transform_points
|
|
from .utils import create_onnx_session
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
class HelloFaceAlignment(object):
|
|
def __init__(self, gpu_id=None):
|
|
expand_ratio = 0.15
|
|
self.face_alignment_net_222 = (
|
|
create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_landmark.onnx'), gpu_id=gpu_id))
|
|
self.onnx_input_name_222 = self.face_alignment_net_222.get_inputs()[0].name
|
|
self.onnx_output_name_222 = [output.name for output in self.face_alignment_net_222.get_outputs()]
|
|
self.face_image_size = 128
|
|
self.face_detector = HelloFaceDet(hf_hub_download('songkey/hello_group_facemodel', filename='hello_face_det.onnx'), gpu_id=gpu_id)
|
|
self.expand_ratio = expand_ratio
|
|
|
|
def onnx_infer(self, input_uint8):
|
|
assert input_uint8.shape[0] == input_uint8.shape[1] == self.face_image_size
|
|
onnx_input = input_uint8.transpose((2, 0, 1)).astype(np.float32)[np.newaxis, :, :, :] / 255.0
|
|
landmark, euler, prob = self.face_alignment_net_222.run(self.onnx_output_name_222,
|
|
{self.onnx_input_name_222: onnx_input})
|
|
|
|
landmark = np.reshape(landmark[0], (2, -1)).transpose((1, 0)) * self.face_image_size
|
|
left_eye_corner = landmark[74]
|
|
right_eye_corner = landmark[96]
|
|
radian = np.arctan2(right_eye_corner[1] - left_eye_corner[1],
|
|
right_eye_corner[0] - left_eye_corner[0] + 0.00000001)
|
|
euler_rad = np.array([euler[0, 0], euler[0, 1], radian], dtype=np.float32)
|
|
prob = prob[0]
|
|
|
|
return landmark, euler_rad, prob
|
|
|
|
def forward(self, src_image, face_box=None, pre_pts=None, iterations=3):
|
|
if pre_pts is None:
|
|
if face_box is None:
|
|
# Detect max size face
|
|
bounding_boxes, _, score = self.face_detector.detect(src_image)
|
|
print("facedet score", score)
|
|
if len(bounding_boxes) == 0:
|
|
return None
|
|
bbox = np.zeros(4, dtype=np.float32)
|
|
if len(bounding_boxes) >= 1:
|
|
max_area = 0.0
|
|
for each_bbox in bounding_boxes:
|
|
area = (each_bbox[2] - each_bbox[0]) * (each_bbox[3] - each_bbox[1])
|
|
if area > max_area:
|
|
bbox[:4] = each_bbox[:4]
|
|
max_area = area
|
|
else:
|
|
bbox = bounding_boxes[0, :4]
|
|
else:
|
|
bbox = face_box.copy()
|
|
M_Face = get_warp_mat_bbox(bbox, 0, self.face_image_size, expand_ratio=self.expand_ratio)
|
|
else:
|
|
left_eye_corner = pre_pts[74]
|
|
right_eye_corner = pre_pts[96]
|
|
|
|
radian = np.arctan2(right_eye_corner[1] - left_eye_corner[1],
|
|
right_eye_corner[0] - left_eye_corner[0] + 0.00000001)
|
|
M_Face = get_warp_mat_bbox_by_gt_pts_float(pre_pts, np.rad2deg(radian), self.face_image_size,
|
|
expand_ratio=self.expand_ratio)
|
|
|
|
face_input = cv2.warpAffine(src_image, M_Face, (self.face_image_size, self.face_image_size))
|
|
landmarks, euler, prob = self.onnx_infer(face_input)
|
|
landmarks = transform_points(landmarks, M_Face, invert=True)
|
|
|
|
# Repeat
|
|
for i in range(iterations - 1):
|
|
M_Face = get_warp_mat_bbox_by_gt_pts_float(landmarks, np.rad2deg(euler[2]), self.face_image_size,
|
|
expand_ratio=self.expand_ratio)
|
|
face_input = cv2.warpAffine(src_image, M_Face, (self.face_image_size, self.face_image_size))
|
|
landmarks, euler, prob = self.onnx_infer(face_input)
|
|
landmarks = transform_points(landmarks, M_Face, invert=True)
|
|
|
|
return_dict = {
|
|
"pt222": landmarks,
|
|
"euler_rad": euler,
|
|
"prob": prob,
|
|
"M_Face": M_Face,
|
|
"face_input": face_input
|
|
}
|
|
|
|
return return_dict
|