Initial code for facial landmark model

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():

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)

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")