Upload src/utils/mp_utils.py with huggingface_hub

src/utils/mp_utils.py

@@ -0,0 +1,95 @@
+import os
+import numpy as np
+import cv2
+import time
+from tqdm import tqdm
+import multiprocessing
+import glob
+
+import mediapipe as mp
+from mediapipe import solutions
+from mediapipe.framework.formats import landmark_pb2
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+from . import face_landmark
+
+CUR_DIR = os.path.dirname(__file__)
+
+
+class LMKExtractor():
+    def __init__(self, FPS=25):
+        # Create an FaceLandmarker object.
+        self.mode = mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE
+        base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/face_landmarker_v2_with_blendshapes.task'))
+        base_options.delegate = mp.tasks.BaseOptions.Delegate.CPU
+        options = vision.FaceLandmarkerOptions(base_options=base_options,
+                                            running_mode=self.mode,
+                                            output_face_blendshapes=True,
+                                            output_facial_transformation_matrixes=True,
+                                            num_faces=1)
+        self.detector = face_landmark.FaceLandmarker.create_from_options(options)
+        self.last_ts = 0
+        self.frame_ms = int(1000 / FPS)
+
+        det_base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/blaze_face_short_range.tflite'))
+        det_options = vision.FaceDetectorOptions(base_options=det_base_options)
+        self.det_detector = vision.FaceDetector.create_from_options(det_options)
+                
+
+    def __call__(self, img):
+        frame = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
+        t0 = time.time()
+        if self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.VIDEO:
+            det_result = self.det_detector.detect(image)
+            if len(det_result.detections) != 1:
+                return None
+            self.last_ts += self.frame_ms
+            try:
+                detection_result, mesh3d = self.detector.detect_for_video(image, timestamp_ms=self.last_ts)
+            except:
+                return None
+        elif self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE:
+            # det_result = self.det_detector.detect(image)
+
+            # if len(det_result.detections) != 1:
+            #     return None
+            try:
+                detection_result, mesh3d = self.detector.detect(image)
+            except:
+                return None
+            
+        
+        bs_list = detection_result.face_blendshapes
+        if len(bs_list) == 1:
+            bs = bs_list[0]
+            bs_values = []
+            for index in range(len(bs)):
+                bs_values.append(bs[index].score)
+            bs_values = bs_values[1:] # remove neutral
+            trans_mat = detection_result.facial_transformation_matrixes[0]
+            face_landmarks_list = detection_result.face_landmarks
+            face_landmarks = face_landmarks_list[0]
+            lmks = []
+            for index in range(len(face_landmarks)):
+                x = face_landmarks[index].x
+                y = face_landmarks[index].y
+                z = face_landmarks[index].z
+                lmks.append([x, y, z])
+            lmks = np.array(lmks)
+            
+            lmks3d = np.array(mesh3d.vertex_buffer)
+            lmks3d = lmks3d.reshape(-1, 5)[:, :3]
+            mp_tris = np.array(mesh3d.index_buffer).reshape(-1, 3) + 1
+
+            return {
+                "lmks": lmks,
+                'lmks3d': lmks3d,
+                "trans_mat": trans_mat,
+                'faces': mp_tris,
+                "bs": bs_values
+            }
+        else:
+            # print('multiple faces in the image: {}'.format(img_path))
+            return None
+