make head posture only use mediapipe results

Rassibassi · Jan 7, 2022 · c49a7c2 · c49a7c2
1 parent 7a4e739
commit c49a7c2
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 14 deletions.
diff --git a/head_posture.py b/head_posture.py
@@ -67,24 +67,40 @@ def main():
                 metric_landmarks, pose_transform_mat = get_metric_landmarks(
                     landmarks.copy(), pcf
                 )
-                model_points = metric_landmarks[0:3, points_idx].T
+
                 image_points = (
                     landmarks[0:2, points_idx].T
                     * np.array([frame_width, frame_height])[None, :]
                 )
 
-                success, rotation_vector, translation_vector = cv2.solvePnP(
-                    model_points,
-                    image_points,
-                    camera_matrix,
-                    dist_coeff,
-                    flags=cv2.cv2.SOLVEPNP_ITERATIVE,
-                )
+                # see here:
+                # https://github.com/google/mediapipe/issues/1379#issuecomment-752534379
+                pose_transform_mat[1:3, :] = -pose_transform_mat[1:3, :]
+                mp_rotation_vector, _ = cv2.Rodrigues(pose_transform_mat[:3, :3])
+                mp_translation_vector = pose_transform_mat[:3, 3, None]
+
+                if False:
+                    # sanity check
+                    # get same result with solvePnP
+                    model_points = metric_landmarks[0:3, points_idx].T
+
+                    success, rotation_vector, translation_vector = cv2.solvePnP(
+                        model_points,
+                        image_points,
+                        camera_matrix,
+                        dist_coeff,
+                        flags=cv2.cv2.SOLVEPNP_ITERATIVE,
+                    )
+
+                    np.testing.assert_almost_equal(mp_rotation_vector, rotation_vector)
+                    np.testing.assert_almost_equal(
+                        mp_translation_vector, translation_vector
+                    )
 
                 (nose_end_point2D, jacobian) = cv2.projectPoints(
                     np.array([(0.0, 0.0, 25.0)]),
-                    rotation_vector,
-                    translation_vector,
+                    mp_rotation_vector,
+                    mp_translation_vector,
                     camera_matrix,
                     dist_coeff,
                 )

diff --git a/videosource.py b/videosource.py
@@ -4,11 +4,12 @@
 
 
 class VideoSource:
-    def __init__(self, display=False, dtype=np.uint8):
+    def __init__(self, flip=False, display=False, dtype=np.uint8):
         self._name = "VideoSource"
         self._capture = None
         self._display = display
         self._dtype = dtype
+        self._flip = flip
 
     def get_fps(self):
         return self._capture.get(cv2.CAP_PROP_FPS)
@@ -32,6 +33,9 @@ def __iter__(self):
     def __next__(self):
         ret, frame = self._capture.read()
 
+        if self._flip:
+            frame = cv2.flip(frame, 3)
+
         if self._display:
             cv2.imshow(f"{self._name} - FPS: {self.get_fps()}", frame)
 
@@ -62,9 +66,10 @@ def __init__(
         fps=30,
         autofocus=0,
         absolute_focus=75,
+        flip=True,
         display=False,
     ):
-        super().__init__(display)
+        super().__init__(flip, display)
         self._capture = cv2.VideoCapture(camera_id)
         self._capture.set(cv2.CAP_PROP_FRAME_WIDTH, width)
         self._capture.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
@@ -75,6 +80,6 @@ def __init__(
 
 
 class FileSource(VideoSource):
-    def __init__(self, file_path, display=False):
-        super().__init__(display)
+    def __init__(self, file_path, flip=False, display=False):
+        super().__init__(flip, display)
         self._capture = cv2.VideoCapture(str(file_path))