From 3a35f3dcf8b07924022fdabcae119624e2483151 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Laure=CE=B7t?= <laurentfainsin@protonmail.com>
Date: Wed, 14 Dec 2022 11:19:04 +0100
Subject: [PATCH] feat: retreive axis from landmarks

---
 src/environment.py | 69 ++++++++++++++++++++++++++++++++++++++++++++--
 src/utils.py       |  9 ++++++
 2 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/src/environment.py b/src/environment.py
index 128cdc0..6332d98 100644
--- a/src/environment.py
+++ b/src/environment.py
@@ -11,6 +11,7 @@ from bodyparts.eye import Eye
 from bodyparts.head import Head
 from bodyparts.moustache import Moustache
 from bodyparts.mouth import Mouth
+from utils import landmark3vec
 
 
 class Environment:
@@ -85,6 +86,9 @@ class Environment:
             # detect keypoints on frame
             self.detect_keypoints()
 
+            # compute face axis
+            self.compute_face_axis()
+
             # draw keypoints on top of fram
             self.draw_keypoints()
 
@@ -115,6 +119,64 @@ class Environment:
             # convert the image back to BGR format
             self.frame = cv2.cvtColor(self.frame, cv2.COLOR_RGB2BGR)
 
+    def compute_face_axis(self) -> None:
+        """Compute the face axis."""
+        if self.results.multi_face_landmarks:
+            for face_landmarks in self.results.multi_face_landmarks:
+                # retreive points
+                left = landmark3vec(face_landmarks.landmark[234], self.screen)
+                right = landmark3vec(face_landmarks.landmark[454], self.screen)
+                bottom = landmark3vec(face_landmarks.landmark[152], self.screen)
+                top = landmark3vec(face_landmarks.landmark[10], self.screen)
+                center = (left + right + bottom + top) / 4
+
+                # compute axis
+                self.x = (right - left) / 2
+                self.y = (top - bottom) / 2
+                self.z = self.x.cross(self.y)
+
+                # normalize axis
+                self.x.normalize_ip()
+                self.y.normalize_ip()
+                self.z.normalize_ip()
+
+                # print horizontal angle to screen
+                self.angle_x = self.x.angle_to(pg.math.Vector3(1, 0, 0))
+                self.angle_y = self.x.angle_to(pg.math.Vector3(0, 0, 1)) - 90
+                self.angle_z = self.y.angle_to(pg.math.Vector3(0, 0, 1)) - 90
+
+                # draw axis on opencv screen
+                cv2.line(
+                    self.frame,
+                    (int(center.x * self.screen.get_width()), int(center.y * self.screen.get_height())),
+                    (
+                        int(center.x * self.screen.get_width() + self.x.x * 100),
+                        int(center.y * self.screen.get_height() + self.x.y * 100),
+                    ),
+                    (0, 0, 255),
+                    2,
+                )
+                cv2.line(
+                    self.frame,
+                    (int(center.x * self.screen.get_width()), int(center.y * self.screen.get_height())),
+                    (
+                        int(center.x * self.screen.get_width() + self.y.x * 100),
+                        int(center.y * self.screen.get_height() + self.y.y * 100),
+                    ),
+                    (0, 255, 0),
+                    2,
+                )
+                cv2.line(
+                    self.frame,
+                    (int(center.x * self.screen.get_width()), int(center.y * self.screen.get_height())),
+                    (
+                        int(center.x * self.screen.get_width() + self.z.x * 100),
+                        int(center.y * self.screen.get_height() + self.z.y * 100),
+                    ),
+                    (255, 0, 0),
+                    2,
+                )
+
     def draw_keypoints(self) -> None:
         """Draw the keypoints on the screen."""
         # draw the face mesh annotations on the image.
@@ -133,7 +195,8 @@ class Environment:
                 # )
 
         # flip the image horizontally for a selfie-view display
-        cv2.imshow("MediaPipe Face Mesh", cv2.flip(self.frame, 1))
+        cv2.imshow("MediaPipe Face Mesh", self.frame)
+        # cv2.imshow("MediaPipe Face Mesh", cv2.flip(self.frame, 1))
 
     def draw_avatar(self) -> None:
         """Draw the avatar on the screen."""
@@ -142,7 +205,7 @@ class Environment:
         for part in self.body_parts.values():
             part.draw(self.screen)
 
-        # self.screen.blit(pg.transform.flip(self.screen, False, False), (0, 0))
-        self.screen.blit(pg.transform.flip(self.screen, True, False), (0, 0))
+        self.screen.blit(self.screen, (0, 0))
+        # self.screen.blit(pg.transform.flip(self.screen, True, False), (0, 0))
 
         pg.display.flip()
diff --git a/src/utils.py b/src/utils.py
index c2a0a1d..0d14595 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -35,3 +35,12 @@ def landmark2vec(landmark, screen):
             screen.get_height(),
         )  # type: ignore
     )
+
+
+def landmark3vec(landmark, screen):
+    """Convert a landmark to a pygame Vector3."""
+    return pg.Vector3(
+        np.clip(landmark.x, 0, 1),
+        np.clip(landmark.y, 0, 1),
+        np.clip(landmark.z, 0, 1),
+    )