File size: 4,925 Bytes
cf812a0 |
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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import math
import time
import cv2
import numpy as np
import onnx
import onnxruntime
def create_onnx_session(onnx_path, providers) -> 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",
# "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
|