Merge pull request roboflow#1316 from David-rn/feat/facemesh_from_med…

…iapipe Facemesh from mediapipe
rolson24 · Jul 5, 2024 · c40734f · c40734f
2 parents e61a107 + 55a9cef
commit c40734f
Show file tree

Hide file tree

Showing 2 changed files with 2,621 additions and 7 deletions.
diff --git a/supervision/keypoint/core.py b/supervision/keypoint/core.py
@@ -243,9 +243,10 @@ def from_mediapipe(
         pose landmark detection inference result.
 
         Args:
-            mediapipe_results (Union[PoseLandmarkerResult, SolutionOutputs]):
-                The output results from Mediapipe. It supports both: the inference
-                result `PoseLandmarker` and the legacy one from `Pose`.
+            mediapipe_results (Union[PoseLandmarkerResult, FaceLandmarkerResult, SolutionOutputs]):
+                The output results from Mediapipe. It support pose and face landmarks
+                from `PoseLandmaker`, `FaceLandmarker` and the legacy ones
+                from `Pose` and `FaceMesh`.
             resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)`
                 representing the resolution of the frame.
 
@@ -283,14 +284,55 @@ def from_mediapipe(
             key_points = sv.KeyPoints.from_mediapipe(
                 pose_landmarker_result, (image_width, image_height))
             ```
+
+            ```python
+            import cv2
+            import mediapipe as mp
+            import supervision as sv
+
+            image = cv2.imread(<SOURCE_IMAGE_PATH>)
+            image_height, image_width, _ = image.shape
+            mediapipe_image = mp.Image(
+                image_format=mp.ImageFormat.SRGB,
+                data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+            options = mp.tasks.vision.FaceLandmarkerOptions(
+                base_options=mp.tasks.BaseOptions(
+                    model_asset_path="face_landmarker.task"
+                ),
+                output_face_blendshapes=True,
+                output_facial_transformation_matrixes=True,
+                num_faces=2)
+
+            FaceLandmarker = mp.tasks.vision.FaceLandmarker
+            with FaceLandmarker.create_from_options(options) as landmarker:
+                face_landmarker_result = landmarker.detect(mediapipe_image)
+
+            key_points = sv.KeyPoints.from_mediapipe(
+                face_landmarker_result, (image_width, image_height))
+            ```
         """  # noqa: E501 // docs
-        results = mediapipe_results.pose_landmarks
-        if not isinstance(mediapipe_results.pose_landmarks, list):
-            if mediapipe_results.pose_landmarks is None:
+        if hasattr(mediapipe_results, "pose_landmarks"):
+            results = mediapipe_results.pose_landmarks
+            if not isinstance(mediapipe_results.pose_landmarks, list):
+                if mediapipe_results.pose_landmarks is None:
+                    results = []
+                else:
+                    results = [
+                        [
+                            landmark
+                            for landmark in mediapipe_results.pose_landmarks.landmark
+                        ]
+                    ]
+        elif hasattr(mediapipe_results, "face_landmarks"):
+            results = mediapipe_results.face_landmarks
+        elif hasattr(mediapipe_results, "multi_face_landmarks"):
+            if mediapipe_results.multi_face_landmarks is None:
                 results = []
             else:
                 results = [
-                    [landmark for landmark in mediapipe_results.pose_landmarks.landmark]
+                    face_landmark.landmark
+                    for face_landmark in mediapipe_results.multi_face_landmarks
                 ]
 
         if len(results) == 0: