From cf5d06729dbb0f677768e20d240c7dfa440f545e Mon Sep 17 00:00:00 2001 From: yakhyo Date: Wed, 23 Apr 2025 17:29:21 +0900 Subject: [PATCH] Initial code for facial landmark model --- uniface/attribute/age_gender.py | 6 +- uniface/face_utils.py | 35 ++++++++++++ uniface/landmark/model.py | 99 +++++++++++++++++++++++++++++++++ 3 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 uniface/landmark/model.py diff --git a/uniface/attribute/age_gender.py b/uniface/attribute/age_gender.py index b51cd7d..9067d63 100644 --- a/uniface/attribute/age_gender.py +++ b/uniface/attribute/age_gender.py @@ -31,7 +31,7 @@ class AgeGender: f"Initializing AgeGender with model={model_name}, " f"input_size={input_size}" ) - + self.input_size = input_size self.input_std = 1.0 self.input_mean = 0.0 @@ -80,7 +80,7 @@ class AgeGender: """ width, height = bbox[2] - bbox[0], bbox[3] - bbox[1] center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 - scale = self.input_size[0] / (max(width, height)*1.5) + scale = self.input_size[0] / (max(width, height) * 1.5) rotation = 0.0 transformed_image, M = bbox_center_alignment(image, center, self.input_size[0], scale, rotation) @@ -117,8 +117,6 @@ class AgeGender: return gender, age - - # TODO: For testing purposes only, remove later def main(): diff --git a/uniface/face_utils.py b/uniface/face_utils.py index 70e81d6..c5bc319 100644 --- a/uniface/face_utils.py +++ b/uniface/face_utils.py @@ -144,3 +144,38 @@ def bbox_center_alignment(image, center, output_size, scale, rotation): cropped = cv2.warpAffine(image, M, (output_size, output_size), borderValue=0.0) return cropped, M + + + +def trans_points2d(pts, M): + new_pts = np.zeros(shape=pts.shape, dtype=np.float32) + for i in range(pts.shape[0]): + pt = pts[i] + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32) + new_pt = np.dot(M, new_pt) + #print('new_pt', new_pt.shape, new_pt) + new_pts[i] = new_pt[0:2] + + return new_pts + + +def trans_points3d(pts, M): + scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1]) + #print(scale) + new_pts = np.zeros(shape=pts.shape, dtype=np.float32) + for i in range(pts.shape[0]): + pt = pts[i] + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32) + new_pt = np.dot(M, new_pt) + #print('new_pt', new_pt.shape, new_pt) + new_pts[i][0:2] = new_pt[0:2] + new_pts[i][2] = pts[i][2] * scale + + return new_pts + + +def trans_points(pts, M): + if pts.shape[1] == 2: + return trans_points2d(pts, M) + else: + return trans_points3d(pts, M) \ No newline at end of file diff --git a/uniface/landmark/model.py b/uniface/landmark/model.py new file mode 100644 index 0000000..40b3ed6 --- /dev/null +++ b/uniface/landmark/model.py @@ -0,0 +1,99 @@ +import cv2 +import onnx +import onnxruntime +import numpy as np + + +# from ..data import get_object + +from uniface.face_utils import bbox_center_alignment, trans_points + +__all__ = [ + 'Landmark', +] + + +class Landmark: + def __init__(self, model_file=None, session=None): + assert model_file is not None + self.model_file = model_file + self.session = session + + model = onnx.load(self.model_file) + + input_mean = 0.0 + input_std = 1.0 + + self.input_mean = input_mean + self.input_std = input_std + # print('input mean and std:', model_file, self.input_mean, self.input_std) + + if self.session is None: + self.session = onnxruntime.InferenceSession(self.model_file, None) + input_cfg = self.session.get_inputs()[0] + input_shape = input_cfg.shape + input_name = input_cfg.name + + self.input_size = tuple(input_shape[2:4][::-1]) + self.input_shape = input_shape + + outputs = self.session.get_outputs() + output_names = [] + for out in outputs: + output_names.append(out.name) + + self.input_name = input_name + self.output_names = output_names + + assert len(self.output_names) == 1 + + output_shape = outputs[0].shape + self.require_pose = False + + self.lmk_dim = 2 + self.lmk_num = output_shape[1]//self.lmk_dim + self.taskname = 'landmark_%dd_%d' % (self.lmk_dim, self.lmk_num) + + def prepare(self, ctx_id, **kwargs): + if ctx_id < 0: + self.session.set_providers(['CPUExecutionProvider']) + + def get(self, img, bbox): + + w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1]) + center = (bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2 + rotate = 0 + _scale = self.input_size[0] / (max(w, h)*1.5) + # print('param:', img.shape, bbox, center, self.input_size, _scale, rotate) + + aimg, M = bbox_center_alignment(img, center, self.input_size[0], _scale, rotate) + input_size = tuple(aimg.shape[0:2][::-1]) + + # assert input_size==self.input_size + blob = cv2.dnn.blobFromImage( + aimg, + 1.0/self.input_std, + input_size, + (self.input_mean, self.input_mean, self.input_mean), + swapRB=True + ) + pred = self.session.run(self.output_names, {self.input_name: blob})[0][0] + if pred.shape[0] >= 3000: + pred = pred.reshape((-1, 3)) + else: + pred = pred.reshape((-1, 2)) + if self.lmk_num < pred.shape[0]: + pred = pred[self.lmk_num*-1:, :] + pred[:, 0:2] += 1 + pred[:, 0:2] *= (self.input_size[0] // 2) + if pred.shape[1] == 3: + pred[:, 2] *= (self.input_size[0] // 2) + + IM = cv2.invertAffineTransform(M) + pred = trans_points(pred, IM) + + return pred + + +if __name__ == "__main__": + model = Landmark("2d106det.onnx")