From 90c436d51ef6fb52f0c492875ffd9e3df62cb634 Mon Sep 17 00:00:00 2001
From: Maurice Van Wassenhove <mauricevanwassenhove@fastmail.com>
Date: Wed, 14 Feb 2024 17:49:33 +0100
Subject: [PATCH] feat: ui fps display

---
 src/app/camera_container.py | 38 +++++++++++++++++++++++++++++++++++++
 src/app/camera_display.py   |  5 +++--
 src/app/main.py             | 24 +++++++++++++----------
 3 files changed, 55 insertions(+), 12 deletions(-)
 create mode 100644 src/app/camera_container.py

diff --git a/src/app/camera_container.py b/src/app/camera_container.py
new file mode 100644
index 0000000..853bc5d
--- /dev/null
+++ b/src/app/camera_container.py
@@ -0,0 +1,38 @@
+from pygame import Rect
+from pygame_gui import UIManager
+from pygame_gui.elements import UILabel
+
+from tracker.mediapipe_pose import MediaPipePose
+
+
+class FPSDisplay(UILabel):
+    """
+    UILabel to display the UI FPS and the camera FPS
+    """
+
+    def __init__(self, ui_manager: UIManager, media_pipe_pose: MediaPipePose):
+        super().__init__(
+            Rect((10, 10), (300, 30)),
+            "UI FPS: -:--  Camera FPS: -:--",
+            ui_manager,
+            anchors={"top": "top", "left": "left"},
+        )
+
+        self.time_deltas = []
+
+        self.media_pipe_pose = media_pipe_pose
+        self.ui_manager = ui_manager
+
+    def update(self, time_delta: float):
+        super().update(time_delta)
+
+        self.time_deltas.append(time_delta)
+        if len(self.time_deltas) > 30:
+            self.time_deltas.pop(0)
+
+        fps = 1 / (sum(self.time_deltas) / len(self.time_deltas))
+
+        ui_fps = fps
+        camera_fps = fps
+
+        self.set_text(f"UI FPS: {ui_fps:.2f}  Camera FPS: {camera_fps:.2f}")
diff --git a/src/app/camera_display.py b/src/app/camera_display.py
index ac26b3b..cc12d8f 100644
--- a/src/app/camera_display.py
+++ b/src/app/camera_display.py
@@ -19,7 +19,7 @@ def __init__(
         media_pipe_pose: MediaPipePose,
     ):
         # Construct the relative rectangle
-        relative_rect = Rect((0, 0), (0, 0))
+        relative_rect = Rect((0, 50), (0, 0))
 
         super().__init__(relative_rect, image_surface, ui_manager)
 
@@ -41,6 +41,7 @@ def __fit_and_center_rect(self):
         :return:
         """
         window_width, window_height = self.ui_manager.window_resolution
+        window_height -= 50  # Subtract the height of the FPSDisplay
         image_width, image_height = self.camera.get_size()
 
         if window_width / window_height > image_width / image_height:
@@ -52,7 +53,7 @@ def __fit_and_center_rect(self):
                 (window_width, window_width * image_height / image_width)
             )
 
-        self.rect.center = (window_width // 2, window_height // 2)
+        self.rect.center = (window_width // 2, window_height // 2 + 50)
 
     def update(self, time_delta: float):
         super().update(time_delta)
diff --git a/src/app/main.py b/src/app/main.py
index d4f3da3..cccde47 100644
--- a/src/app/main.py
+++ b/src/app/main.py
@@ -2,6 +2,7 @@
 import pygame.camera
 from pygame_gui import UIManager
 
+from app.camera_container import FPSDisplay
 from app.camera_display import CameraDisplay
 from tracker.mediapipe_pose import MediaPipePose
 
@@ -20,15 +21,17 @@ def main():
 
     media_pipe_pose = MediaPipePose()
 
-    num_connected_cameras = 1
-    cam_names = pygame.camera.list_cameras()
-    for cam_name in cam_names[:num_connected_cameras]:
-        CameraDisplay(
-            camera_id=cam_name,
-            image_surface=window_surface,
-            ui_manager=ui_manager,
-            media_pipe_pose=media_pipe_pose,
-        )
+    FPSDisplay(
+        ui_manager=ui_manager,
+        media_pipe_pose=media_pipe_pose,
+    )
+
+    CameraDisplay(
+        image_surface=pygame.Surface((640, 480)),
+        camera_id="/dev/video0",
+        ui_manager=ui_manager,
+        media_pipe_pose=media_pipe_pose,
+    )
 
     clock = pygame.time.Clock()
     is_running = True
@@ -41,12 +44,13 @@ def main():
                     is_running = False
                 case pygame.VIDEORESIZE:
                     ui_manager.set_window_resolution(window_surface.get_size())
-                    window_surface.fill(pygame.Color("#000000"))
 
             ui_manager.process_events(event)
 
         ui_manager.update(time_delta)
+        window_surface.fill(pygame.Color("#000000"))
         ui_manager.draw_ui(window_surface)
+
         pygame.display.update()