init
This commit is contained in:
songkey-win
@@ -0,0 +1,13 @@
|
||||
# coding: utf-8
|
||||
|
||||
# @File : __init__.py
|
||||
# @Author : Songkey
|
||||
# @Email : songkey@pku.edu.cn
|
||||
# @Date : 8/28/2024
|
||||
# @Desc :
|
||||
|
||||
from .hello_arkit import HelloARKitBSPred
|
||||
from .hello_face_det import HelloFaceDet
|
||||
from .hello_camera_demo import HelloCameraDemo
|
||||
from .hello_3dmm import Hello3DMMPred
|
||||
from .hello_face_alignment import HelloFaceAlignment
|
||||
@@ -0,0 +1,140 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
@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
|
||||
import cv2
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
from .utils import get_warp_mat_bbox_by_gt_pts_float, create_onnx_session
|
||||
|
||||
def crop_transl_to_full_transl(crop_trans, crop_center, scale, full_center, focal_length):
|
||||
"""
|
||||
:param crop_trans: (3), float
|
||||
:param crop_center: (2), float
|
||||
:param scale: (1), float
|
||||
:param full_center: (2), float
|
||||
:param focal_length: (1), float
|
||||
:return:
|
||||
"""
|
||||
crop_c_x, crop_c_y = crop_center
|
||||
full_c_x, full_c_y = full_center
|
||||
bs = 2 * focal_length / scale / crop_trans[2]
|
||||
full_x = crop_trans[0] - 2 * (crop_c_x - full_c_x) / bs
|
||||
full_y = crop_trans[1] + 2 * (crop_c_y - full_c_y) / bs
|
||||
full_z = crop_trans[2] * scale
|
||||
|
||||
full_trans = np.array([full_x, full_y, full_z], dtype=np.float32)
|
||||
|
||||
return full_trans
|
||||
|
||||
class Hello3DMMPred(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.deep3d_pred_net = create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_3dmm.onnx'), gpu_id=gpu_id)
|
||||
self.deep3d_pred_net_input_name = self.deep3d_pred_net.get_inputs()[0].name
|
||||
self.deep3d_pred_net_output_name = [output.name for output in self.deep3d_pred_net.get_outputs()]
|
||||
|
||||
self.image_size = 224
|
||||
self.camera_init_z = -0.4
|
||||
self.camera_init_focal_len = 386.2879122887948
|
||||
self.used_focal_len = -5.0 / self.camera_init_z * self.camera_init_focal_len
|
||||
self.id_dims = 526
|
||||
self.exp_dims = 203
|
||||
self.tex_dims = 439
|
||||
|
||||
def forward_params(self, src_image, src_pt):
|
||||
align_mat_info = get_warp_mat_bbox_by_gt_pts_float(src_pt, base_angle=0, dst_size=self.image_size, expand_ratio=0.35, return_info=True)
|
||||
align_mat = align_mat_info["M"]
|
||||
|
||||
align_image_rgb_uint8 = cv2.cvtColor(cv2.warpAffine(src_image, align_mat, (self.image_size, self.image_size)), cv2.COLOR_BGR2RGB)
|
||||
|
||||
# cv2.imshow('align_image_rgb_uint8', align_image_rgb_uint8)
|
||||
|
||||
align_image_rgb_fp32 = align_image_rgb_uint8.astype(np.float32) / 255.0
|
||||
align_image_rgb_fp32_onnx_input = align_image_rgb_fp32.copy().transpose((2, 0, 1))[np.newaxis, ...]
|
||||
pred_coeffs = self.deep3d_pred_net.run(self.deep3d_pred_net_output_name,
|
||||
{self.deep3d_pred_net_input_name: align_image_rgb_fp32_onnx_input})[0]
|
||||
|
||||
angles = pred_coeffs[:, self.id_dims + self.exp_dims + self.tex_dims:self.id_dims + self.exp_dims + self.tex_dims + 3]
|
||||
translations = pred_coeffs[:, self.id_dims + self.exp_dims + self.tex_dims + 3 + 27:]
|
||||
|
||||
crop_global_transl = crop_transl_to_full_transl(translations[0],
|
||||
crop_center=[align_mat_info["center_x"],
|
||||
align_mat_info["center_y"]],
|
||||
scale=align_mat_info["scale"],
|
||||
full_center=[src_image.shape[1] * 0.5, src_image.shape[0] * 0.5],
|
||||
focal_length=self.used_focal_len)
|
||||
return angles, crop_global_transl[np.newaxis, :]
|
||||
|
||||
def compute_rotation_matrix(self, angles):
|
||||
n_b = angles.shape[0]
|
||||
sinx = np.sin(angles[:, 0])
|
||||
siny = np.sin(angles[:, 1])
|
||||
sinz = np.sin(angles[:, 2])
|
||||
cosx = np.cos(angles[:, 0])
|
||||
cosy = np.cos(angles[:, 1])
|
||||
cosz = np.cos(angles[:, 2])
|
||||
rotXYZ = np.eye(3).reshape(1, 3, 3).repeat(n_b*3, 0).reshape(3, n_b, 3, 3)
|
||||
rotXYZ[0, :, 1, 1] = cosx
|
||||
rotXYZ[0, :, 1, 2] = -sinx
|
||||
rotXYZ[0, :, 2, 1] = sinx
|
||||
rotXYZ[0, :, 2, 2] = cosx
|
||||
rotXYZ[1, :, 0, 0] = cosy
|
||||
rotXYZ[1, :, 0, 2] = siny
|
||||
rotXYZ[1, :, 2, 0] = -siny
|
||||
rotXYZ[1, :, 2, 2] = cosy
|
||||
rotXYZ[2, :, 0, 0] = cosz
|
||||
rotXYZ[2, :, 0, 1] = -sinz
|
||||
rotXYZ[2, :, 1, 0] = sinz
|
||||
rotXYZ[2, :, 1, 1] = cosz
|
||||
rotation = np.matmul(np.matmul(rotXYZ[2], rotXYZ[1]), rotXYZ[0])
|
||||
return rotation.transpose(0, 2, 1)
|
||||
|
||||
def rigid_transform(self, vs, rot, trans):
|
||||
vs_r = np.matmul(vs, rot)
|
||||
vs_t = vs_r + trans.reshape(-1, 1, 3)
|
||||
return vs_t
|
||||
|
||||
def perspective_projection_points(self, points, image_w, image_h, used_focal_len):
|
||||
batch_size = points.shape[0]
|
||||
K = np.zeros([batch_size, 3, 3])
|
||||
K[:, 0, 0] = used_focal_len
|
||||
K[:, 1, 1] = used_focal_len
|
||||
K[:, 2, 2] = 1.
|
||||
K[:, 0, 2] = image_w * 0.5
|
||||
K[:, 1, 2] = image_h * 0.5
|
||||
|
||||
reverse_z = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])[np.newaxis, :, :].repeat(batch_size, 0)
|
||||
|
||||
# Transform points
|
||||
aug_projection = np.matmul(points, reverse_z)
|
||||
aug_projection = np.matmul(aug_projection, K.transpose((0, 2, 1)))
|
||||
|
||||
# Apply perspective distortion
|
||||
projected_points = aug_projection[:, :, :2] / aug_projection[:, :, 2:]
|
||||
return projected_points
|
||||
|
||||
def get_project_points_rect(self, angle, trans, image_w, image_h):
|
||||
vs = np.array(
|
||||
[[-1, -1, 0], [-1, 1, 0], [1, 1, 0], [1, -1, 0]],
|
||||
) * 0.05
|
||||
vs = vs[np.newaxis, :, :]
|
||||
|
||||
rotation = self.compute_rotation_matrix(angle)
|
||||
translation = trans.copy()
|
||||
translation[0, 2] *= 0.05
|
||||
|
||||
vs_t = self.rigid_transform(vs, rotation, translation)
|
||||
|
||||
project_points = self.perspective_projection_points(vs_t, image_w, image_h, self.used_focal_len*0.05)
|
||||
project_points = np.stack([project_points[:, :, 0], image_h - project_points[:, :, 1]], axis=2)
|
||||
|
||||
return project_points[0]
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
@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
|
||||
import cv2
|
||||
from huggingface_hub import hf_hub_download
|
||||
from .utils import create_onnx_session, get_warp_mat_bbox_by_gt_pts_float
|
||||
|
||||
class HelloARKitBSPred(object):
|
||||
def __init__(self, gpu_id=0):
|
||||
self.face_rig_net = create_onnx_session(hf_hub_download('songkey/hello_group_facemodel', filename='hello_arkit_blendshape.onnx'), gpu_id=gpu_id)
|
||||
self.onnx_input_name = self.face_rig_net.get_inputs()[0].name
|
||||
self.onnx_output_name = [output.name for output in self.face_rig_net.get_outputs()]
|
||||
self.image_size = 224
|
||||
self.expand_ratio = 0.15
|
||||
|
||||
def forward(self, src_image, src_pt):
|
||||
left_eye_corner = src_pt[74]
|
||||
right_eye_corner = src_pt[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)
|
||||
align_warp_mat = get_warp_mat_bbox_by_gt_pts_float(src_pt, base_angle=rotate_angle, dst_size=self.image_size,
|
||||
expand_ratio=self.expand_ratio)
|
||||
face_rig_input = cv2.warpAffine(src_image, align_warp_mat, (self.image_size, self.image_size))
|
||||
|
||||
face_rig_onnx_input = face_rig_input.transpose((2, 0, 1)).astype(np.float32)[np.newaxis, :, :, :] / 255.0
|
||||
face_rig_params = self.face_rig_net.run(self.onnx_output_name,
|
||||
{self.onnx_input_name: face_rig_onnx_input})
|
||||
face_rig_params = face_rig_params[0][0]
|
||||
return face_rig_params
|
||||
@@ -0,0 +1,377 @@
|
||||
"""
|
||||
@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.01
|
||||
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
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
@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
|
||||
@@ -0,0 +1,326 @@
|
||||
"""
|
||||
@File : test.py
|
||||
@Author : Songkey
|
||||
@Email : songkey@pku.edu.cn
|
||||
@Date : 11/1/2024
|
||||
@Desc : Created by Zemin An (an.zemin@hellogroup.com)
|
||||
"""
|
||||
|
||||
from abc import ABCMeta, abstractmethod
|
||||
import cv2
|
||||
import numpy as np
|
||||
from scipy.special import softmax
|
||||
import os.path as osp
|
||||
from .utils import create_onnx_session
|
||||
|
||||
songkey_weights_dir = 'pretrained_models'
|
||||
|
||||
_COLORS = (
|
||||
np.array(
|
||||
[
|
||||
0.000,
|
||||
0.447,
|
||||
0.741,
|
||||
]
|
||||
)
|
||||
.astype(np.float32)
|
||||
.reshape(-1, 3)
|
||||
)
|
||||
|
||||
def get_resize_matrix(raw_shape, dst_shape, keep_ratio):
|
||||
"""
|
||||
Get resize matrix for resizing raw img to input size
|
||||
:param raw_shape: (width, height) of raw image
|
||||
:param dst_shape: (width, height) of input image
|
||||
:param keep_ratio: whether keep original ratio
|
||||
:return: 3x3 Matrix
|
||||
"""
|
||||
r_w, r_h = raw_shape
|
||||
d_w, d_h = dst_shape
|
||||
Rs = np.eye(3)
|
||||
if keep_ratio:
|
||||
C = np.eye(3)
|
||||
C[0, 2] = -r_w / 2
|
||||
C[1, 2] = -r_h / 2
|
||||
|
||||
if r_w / r_h < d_w / d_h:
|
||||
ratio = d_h / r_h
|
||||
else:
|
||||
ratio = d_w / r_w
|
||||
Rs[0, 0] *= ratio
|
||||
Rs[1, 1] *= ratio
|
||||
|
||||
T = np.eye(3)
|
||||
T[0, 2] = 0.5 * d_w
|
||||
T[1, 2] = 0.5 * d_h
|
||||
return T @ Rs @ C
|
||||
else:
|
||||
Rs[0, 0] *= d_w / r_w
|
||||
Rs[1, 1] *= d_h / r_h
|
||||
return Rs
|
||||
|
||||
def warp_boxes(boxes, M, width, height):
|
||||
"""Apply transform to boxes
|
||||
Copy from nanodet/data/transform/warp.py
|
||||
"""
|
||||
n = len(boxes)
|
||||
if n:
|
||||
# warp points
|
||||
xy = np.ones((n * 4, 3))
|
||||
xy[:, :2] = boxes[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(
|
||||
n * 4, 2
|
||||
) # x1y1, x2y2, x1y2, x2y1
|
||||
xy = xy @ M.T # transform
|
||||
xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8) # rescale
|
||||
# create new boxes
|
||||
x = xy[:, [0, 2, 4, 6]]
|
||||
y = xy[:, [1, 3, 5, 7]]
|
||||
xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
|
||||
# clip boxes
|
||||
xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)
|
||||
xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)
|
||||
return xy.astype(np.float32)
|
||||
else:
|
||||
return boxes
|
||||
|
||||
def overlay_bbox_cv(img, all_box, class_names):
|
||||
"""Draw result boxes
|
||||
Copy from nanodet/util/visualization.py
|
||||
"""
|
||||
# all_box array of [label, x0, y0, x1, y1, score]
|
||||
all_box.sort(key=lambda v: v[5])
|
||||
for box in all_box:
|
||||
label, x0, y0, x1, y1, score = box
|
||||
# color = self.cmap(i)[:3]
|
||||
color = (_COLORS[label] * 255).astype(np.uint8).tolist()
|
||||
text = "{}:{:.1f}%".format(class_names[label], score * 100)
|
||||
txt_color = (0, 0, 0) if np.mean(_COLORS[label]) > 0.5 else (255, 255, 255)
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
txt_size = cv2.getTextSize(text, font, 0.5, 2)[0]
|
||||
cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)
|
||||
|
||||
cv2.rectangle(
|
||||
img,
|
||||
(x0, y0 - txt_size[1] - 1),
|
||||
(x0 + txt_size[0] + txt_size[1], y0 - 1),
|
||||
color,
|
||||
-1,
|
||||
)
|
||||
cv2.putText(img, text, (x0, y0 - 1), font, 0.5, txt_color, thickness=1)
|
||||
return img
|
||||
|
||||
def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
|
||||
"""
|
||||
|
||||
Args:
|
||||
box_scores (N, 5): boxes in corner-form and probabilities.
|
||||
iou_threshold: intersection over union threshold.
|
||||
top_k: keep top_k results. If k <= 0, keep all the results.
|
||||
candidate_size: only consider the candidates with the highest scores.
|
||||
Returns:
|
||||
picked: a list of indexes of the kept boxes
|
||||
"""
|
||||
scores = box_scores[:, -1]
|
||||
boxes = box_scores[:, :-1]
|
||||
picked = []
|
||||
# _, indexes = scores.sort(descending=True)
|
||||
indexes = np.argsort(scores)
|
||||
# indexes = indexes[:candidate_size]
|
||||
indexes = indexes[-candidate_size:]
|
||||
while len(indexes) > 0:
|
||||
# current = indexes[0]
|
||||
current = indexes[-1]
|
||||
picked.append(current)
|
||||
if 0 < top_k == len(picked) or len(indexes) == 1:
|
||||
break
|
||||
current_box = boxes[current, :]
|
||||
# indexes = indexes[1:]
|
||||
indexes = indexes[:-1]
|
||||
rest_boxes = boxes[indexes, :]
|
||||
iou = iou_of(
|
||||
rest_boxes,
|
||||
np.expand_dims(current_box, axis=0),
|
||||
)
|
||||
indexes = indexes[iou <= iou_threshold]
|
||||
|
||||
return box_scores[picked, :]
|
||||
|
||||
|
||||
def iou_of(boxes0, boxes1, eps=1e-5):
|
||||
"""Return intersection-over-union (Jaccard index) of boxes.
|
||||
|
||||
Args:
|
||||
boxes0 (N, 4): ground truth boxes.
|
||||
boxes1 (N or 1, 4): predicted boxes.
|
||||
eps: a small number to avoid 0 as denominator.
|
||||
Returns:
|
||||
iou (N): IoU values.
|
||||
"""
|
||||
overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
|
||||
overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])
|
||||
|
||||
overlap_area = area_of(overlap_left_top, overlap_right_bottom)
|
||||
area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
|
||||
area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
|
||||
return overlap_area / (area0 + area1 - overlap_area + eps)
|
||||
|
||||
|
||||
def area_of(left_top, right_bottom):
|
||||
"""Compute the areas of rectangles given two corners.
|
||||
|
||||
Args:
|
||||
left_top (N, 2): left top corner.
|
||||
right_bottom (N, 2): right bottom corner.
|
||||
|
||||
Returns:
|
||||
area (N): return the area.
|
||||
"""
|
||||
hw = np.clip(right_bottom - left_top, 0.0, None)
|
||||
return hw[..., 0] * hw[..., 1]
|
||||
|
||||
|
||||
class NanoDetABC(metaclass=ABCMeta):
|
||||
def __init__(
|
||||
self,
|
||||
input_shape=[272, 160],
|
||||
reg_max=7,
|
||||
strides=[8, 16, 32],
|
||||
prob_threshold=0.4,
|
||||
iou_threshold=0.3,
|
||||
num_candidate=1000,
|
||||
top_k=-1,
|
||||
class_names=["face"]
|
||||
):
|
||||
self.strides = strides
|
||||
self.input_shape = input_shape
|
||||
self.reg_max = reg_max
|
||||
self.prob_threshold = prob_threshold
|
||||
self.iou_threshold = iou_threshold
|
||||
self.num_candidate = num_candidate
|
||||
self.top_k = top_k
|
||||
self.img_mean = [103.53, 116.28, 123.675]
|
||||
self.img_std = [57.375, 57.12, 58.395]
|
||||
self.input_size = (self.input_shape[1], self.input_shape[0])
|
||||
self.class_names = class_names
|
||||
self.num_classes = len(self.class_names)
|
||||
|
||||
def preprocess(self, img):
|
||||
# resize image
|
||||
ResizeM = get_resize_matrix((img.shape[1], img.shape[0]), self.input_size, True)
|
||||
img_resize = cv2.warpPerspective(img, ResizeM, dsize=self.input_size)
|
||||
|
||||
# normalize image
|
||||
img_input = img_resize.astype(np.float32) / 255
|
||||
img_mean = np.array(self.img_mean, dtype=np.float32).reshape(1, 1, 3) / 255
|
||||
img_std = np.array(self.img_std, dtype=np.float32).reshape(1, 1, 3) / 255
|
||||
img_input = (img_input - img_mean) / img_std
|
||||
|
||||
# expand dims
|
||||
img_input = np.transpose(img_input, [2, 0, 1])
|
||||
img_input = np.expand_dims(img_input, axis=0)
|
||||
return img_input, ResizeM
|
||||
|
||||
def postprocess(self, scores, raw_boxes, ResizeM, raw_shape):
|
||||
# generate centers
|
||||
decode_boxes = []
|
||||
select_scores = []
|
||||
for stride, box_distribute, score in zip(self.strides, raw_boxes, scores):
|
||||
# centers
|
||||
fm_h = self.input_shape[0] / stride
|
||||
fm_w = self.input_shape[1] / stride
|
||||
|
||||
h_range = np.arange(fm_h)
|
||||
w_range = np.arange(fm_w)
|
||||
ww, hh = np.meshgrid(w_range, h_range)
|
||||
|
||||
ct_row = hh.flatten() * stride
|
||||
ct_col = ww.flatten() * stride
|
||||
|
||||
center = np.stack((ct_col, ct_row, ct_col, ct_row), axis=1)
|
||||
|
||||
# box distribution to distance
|
||||
reg_range = np.arange(self.reg_max + 1)
|
||||
box_distance = box_distribute.reshape((-1, self.reg_max + 1))
|
||||
box_distance = softmax(box_distance, axis=1)
|
||||
box_distance = box_distance * np.expand_dims(reg_range, axis=0)
|
||||
box_distance = np.sum(box_distance, axis=1).reshape((-1, 4))
|
||||
box_distance = box_distance * stride
|
||||
|
||||
# top K candidate
|
||||
topk_idx = np.argsort(score.max(axis=1))[::-1]
|
||||
topk_idx = topk_idx[: self.num_candidate]
|
||||
center = center[topk_idx]
|
||||
score = score[topk_idx]
|
||||
box_distance = box_distance[topk_idx]
|
||||
|
||||
# decode box
|
||||
decode_box = center + [-1, -1, 1, 1] * box_distance
|
||||
|
||||
select_scores.append(score)
|
||||
decode_boxes.append(decode_box)
|
||||
|
||||
# nms
|
||||
bboxes = np.concatenate(decode_boxes, axis=0)
|
||||
confidences = np.concatenate(select_scores, axis=0)
|
||||
picked_box_probs = []
|
||||
picked_labels = []
|
||||
for class_index in range(0, confidences.shape[1]):
|
||||
probs = confidences[:, class_index]
|
||||
mask = probs > self.prob_threshold
|
||||
probs = probs[mask]
|
||||
if probs.shape[0] == 0:
|
||||
continue
|
||||
subset_boxes = bboxes[mask, :]
|
||||
box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1)
|
||||
box_probs = hard_nms(
|
||||
box_probs,
|
||||
iou_threshold=self.iou_threshold,
|
||||
top_k=self.top_k,
|
||||
)
|
||||
picked_box_probs.append(box_probs)
|
||||
picked_labels.extend([class_index] * box_probs.shape[0])
|
||||
if not picked_box_probs:
|
||||
return np.array([]), np.array([]), np.array([])
|
||||
picked_box_probs = np.concatenate(picked_box_probs)
|
||||
|
||||
# resize output boxes
|
||||
picked_box_probs[:, :4] = warp_boxes(
|
||||
picked_box_probs[:, :4], np.linalg.inv(ResizeM), raw_shape[1], raw_shape[0]
|
||||
)
|
||||
return (
|
||||
picked_box_probs[:, :4].astype(np.int32),
|
||||
np.array(picked_labels),
|
||||
picked_box_probs[:, 4],
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def infer_image(self, img_input):
|
||||
pass
|
||||
|
||||
def detect(self, img):
|
||||
raw_shape = img.shape
|
||||
img_input, ResizeM = self.preprocess(img)
|
||||
scores, raw_boxes = self.infer_image(img_input)
|
||||
if scores[0].ndim == 1: # handling num_classes=1 case
|
||||
scores = [x[:, None] for x in scores]
|
||||
bbox, label, score = self.postprocess(scores, raw_boxes, ResizeM, raw_shape)
|
||||
|
||||
return bbox, label, score
|
||||
|
||||
class HelloFaceDet(NanoDetABC):
|
||||
def __init__(self, model_path=osp.join(songkey_weights_dir, 'face/nanodet_humandet_320-192_220302_model_20220315_test3.onnx'), gpu_id=None, *args, **kwargs):
|
||||
super(HelloFaceDet, self).__init__(*args, **kwargs)
|
||||
# print("Using ONNX as inference backend")
|
||||
# print(f"Using weight: {model_path}")
|
||||
|
||||
# load model
|
||||
self.model_path = model_path
|
||||
self.ort_session = create_onnx_session(model_path, gpu_id=gpu_id)
|
||||
self.input_name = self.ort_session.get_inputs()[0].name
|
||||
|
||||
def infer_image(self, img_input):
|
||||
inference_results = self.ort_session.run(None, {self.input_name: img_input})
|
||||
|
||||
scores = [np.squeeze(x) for x in inference_results[:3]]
|
||||
raw_boxes = [np.squeeze(x) for x in inference_results[3:]]
|
||||
return scores, raw_boxes
|
||||
@@ -0,0 +1,118 @@
|
||||
# coding: utf-8
|
||||
|
||||
# @File : utils.py
|
||||
# @Author : Songkey
|
||||
# @Email : songkey@pku.edu.cn
|
||||
# @Date : 8/18/2024
|
||||
# @Desc :
|
||||
|
||||
import onnx, onnxruntime
|
||||
import time
|
||||
import cv2
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
def create_onnx_session(onnx_path, gpu_id=None)->onnxruntime.InferenceSession:
|
||||
start = time.perf_counter()
|
||||
onnx_model = onnx.load(onnx_path)
|
||||
onnx.checker.check_model(onnx_model)
|
||||
providers = [
|
||||
('CUDAExecutionProvider', {
|
||||
'device_id': int(gpu_id),
|
||||
'arena_extend_strategy': 'kNextPowerOfTwo',
|
||||
#'cuda_mem_limit': 5 * 1024 * 1024 * 1024,
|
||||
'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
||||
'do_copy_in_default_stream': True,
|
||||
}),
|
||||
'CPUExecutionProvider',
|
||||
] if (gpu_id is not None and gpu_id >= 0) else ['CPUExecutionProvider']
|
||||
|
||||
sess = onnxruntime.InferenceSession(onnx_path, providers=providers)
|
||||
print('create onnx session cost: {:.3f}s. {}'.format(time.perf_counter() - start, onnx_path))
|
||||
return sess
|
||||
|
||||
def smoothing_factor(t_e, cutoff):
|
||||
r = 2 * math.pi * cutoff * t_e
|
||||
return r / (r + 1)
|
||||
|
||||
def exponential_smoothing(a, x, x_prev):
|
||||
return a * x + (1 - a) * x_prev
|
||||
|
||||
class OneEuroFilter:
|
||||
def __init__(self, dx0=0.0, d_cutoff=1.0):
|
||||
"""Initialize the one euro filter."""
|
||||
# self.min_cutoff = float(min_cutoff)
|
||||
# self.beta = float(beta)
|
||||
self.d_cutoff = float(d_cutoff)
|
||||
self.dx_prev = float(dx0)
|
||||
# self.t_e = fcmin
|
||||
|
||||
def __call__(self, x, x_prev, fcmin=1.0, min_cutoff=1.0, beta=0.0):
|
||||
if x_prev is None:
|
||||
return x
|
||||
# t_e = 1
|
||||
a_d = smoothing_factor(fcmin, self.d_cutoff)
|
||||
dx = (x - x_prev) / fcmin
|
||||
dx_hat = exponential_smoothing(a_d, dx, self.dx_prev)
|
||||
cutoff = min_cutoff + beta * abs(dx_hat)
|
||||
a = smoothing_factor(fcmin, cutoff)
|
||||
x_hat = exponential_smoothing(a, x, x_prev)
|
||||
self.dx_prev = dx_hat
|
||||
return x_hat
|
||||
|
||||
def get_warp_mat_bbox(face_bbox, base_angle, dst_size=128, expand_ratio=0.15, aug_angle=0.0, aug_scale=1.0):
|
||||
face_x_min, face_y_min, face_x_max, face_y_max = face_bbox
|
||||
face_x_center = (face_x_min + face_x_max) / 2
|
||||
face_y_center = (face_y_min + face_y_max) / 2
|
||||
face_width = face_x_max - face_x_min
|
||||
face_height = face_y_max - face_y_min
|
||||
scale = dst_size / max(face_width, face_height) * (1 - expand_ratio) * aug_scale
|
||||
M = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle + aug_angle, scale=scale)
|
||||
offset = [dst_size / 2 - face_x_center, dst_size / 2 - face_y_center]
|
||||
M[:, 2] += offset
|
||||
return M
|
||||
|
||||
def transform_points(points, mat, invert=False):
|
||||
if invert:
|
||||
mat = cv2.invertAffineTransform(mat)
|
||||
points = np.expand_dims(points, axis=1)
|
||||
points = cv2.transform(points, mat, points.shape)
|
||||
points = np.squeeze(points)
|
||||
return points
|
||||
|
||||
def get_warp_mat_bbox_by_gt_pts_float(gt_pts, base_angle=0.0, dst_size=128, expand_ratio=0.15, return_info=False):
|
||||
# step 1
|
||||
face_x_min, face_x_max = np.min(gt_pts[:, 0]), np.max(gt_pts[:, 0])
|
||||
face_y_min, face_y_max = np.min(gt_pts[:, 1]), np.max(gt_pts[:, 1])
|
||||
face_x_center = (face_x_min + face_x_max) / 2
|
||||
face_y_center = (face_y_min + face_y_max) / 2
|
||||
M_step_1 = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle, scale=1.0)
|
||||
pts_step_1 = transform_points(gt_pts, M_step_1)
|
||||
face_x_min_step_1, face_x_max_step_1 = np.min(pts_step_1[:, 0]), np.max(pts_step_1[:, 0])
|
||||
face_y_min_step_1, face_y_max_step_1 = np.min(pts_step_1[:, 1]), np.max(pts_step_1[:, 1])
|
||||
# step 2
|
||||
face_width = face_x_max_step_1 - face_x_min_step_1
|
||||
face_height = face_y_max_step_1 - face_y_min_step_1
|
||||
scale = dst_size / max(face_width, face_height) * (1 - expand_ratio)
|
||||
M_step_2 = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle, scale=scale)
|
||||
pts_step_2 = transform_points(gt_pts, M_step_2)
|
||||
face_x_min_step_2, face_x_max_step_2 = np.min(pts_step_2[:, 0]), np.max(pts_step_2[:, 0])
|
||||
face_y_min_step_2, face_y_max_step_2 = np.min(pts_step_2[:, 1]), np.max(pts_step_2[:, 1])
|
||||
face_x_center_step_2 = (face_x_min_step_2 + face_x_max_step_2) / 2
|
||||
face_y_center_step_2 = (face_y_min_step_2 + face_y_max_step_2) / 2
|
||||
|
||||
M = cv2.getRotationMatrix2D((face_x_center, face_y_center), angle=base_angle, scale=scale)
|
||||
offset = [dst_size / 2 - face_x_center_step_2, dst_size / 2 - face_y_center_step_2]
|
||||
M[:, 2] += offset
|
||||
|
||||
if not return_info:
|
||||
return M
|
||||
else:
|
||||
transform_info = {
|
||||
"M": M,
|
||||
"center_x": face_x_center,
|
||||
"center_y": face_y_center,
|
||||
"rotate_angle": base_angle,
|
||||
"scale": scale
|
||||
}
|
||||
return transform_info
|
||||
Reference in New Issue
Block a user