From cffaae09f20eb4220c617a42835d3d5c82cf02ef Mon Sep 17 00:00:00 2001 From: Maurice Van Wassenhove Date: Tue, 20 Feb 2024 14:46:16 +0100 Subject: [PATCH] feat: multiprocessing for graphs --- src/app/latency_graph.py | 99 +++++++++++++++++++++++--------- src/app/main.py | 50 +++++++++------- src/app/marker_location_graph.py | 7 +++ 3 files changed, 108 insertions(+), 48 deletions(-) create mode 100644 src/app/marker_location_graph.py diff --git a/src/app/latency_graph.py b/src/app/latency_graph.py index ac1b4cb..b40bcda 100644 --- a/src/app/latency_graph.py +++ b/src/app/latency_graph.py @@ -1,6 +1,44 @@ +import io +from multiprocessing import Pool + import pygame from pygame_gui.elements import UIImage +from tracker.mediapipe_pose import MediaPipePose + + +def plot_latency_graph( + ui_latencies: [float], mediapipe_latencies: [float], frame: int +) -> io.BytesIO: + """ + Plots the latency graph using matplotlib writing the plot to a BytesIO object + Cannot be a method of the LatencyGraph class because it is used in a multiprocessing Pool + :param ui_latencies: The UI latencies to plot + :param mediapipe_latencies: The mediapipe inference latency to plot + """ + import matplotlib.pyplot as plt + import io + import numpy as np + + fig, ax = plt.subplots() + + # generate x values + x = np.arange(frame - len(ui_latencies), frame) + # Plot ui latencies in blue and mediapipe latencies in orange + ax.plot(x, ui_latencies, label="UI Latency", color="blue") + ax.plot(x, mediapipe_latencies, label="MediaPipe Latency", color="orange") + + ax.set_xlabel("Frame") + ax.set_ylabel("Latency (ms)") + ax.set_title("Latency Graph") + + buf = io.BytesIO() + fig.savefig(buf) + buf.seek(0) + plt.close(fig) + + return buf + class LatencyGraph(UIImage): """ @@ -8,37 +46,46 @@ class LatencyGraph(UIImage): """ def __init__( - self, relative_rect: pygame.Rect, image_surface: pygame.surface.Surface + self, + relative_rect: pygame.Rect, + image_surface: pygame.surface.Surface, + pool: Pool, + media_pipe_pose: MediaPipePose, ): super().__init__(relative_rect, image_surface) - self.latencies = [] - - def plot_latency_graph(self) -> pygame.surface.Surface: - """ - Plots the latency graph using matplotlib and converts the plot to a pygame surface - """ - import matplotlib.pyplot as plt - import io - import numpy as np - fig, ax = plt.subplots() - ax.plot(np.arange(len(self.latencies)), self.latencies) - ax.set_xlabel("Frames") - ax.set_ylabel("Latency (s)") - ax.set_title("Latency Graph") - - buf = io.BytesIO() - fig.savefig(buf) - buf.seek(0) - plt.close(fig) - - return pygame.image.load(buf) + self.ui_latencies = [] + self.mediapipe_latencies = [] + self.pool = pool + self.image = image_surface + self.plot_async_result = None + self.media_pipe_pose = media_pipe_pose + self.frame = 0 def update(self, time_delta: float): super().update(time_delta) - self.latencies.append(time_delta) - if len(self.latencies) > 30: - self.latencies.pop(0) + self.frame += 1 + + self.ui_latencies.append(time_delta) + if len(self.ui_latencies) > 30: + self.ui_latencies.pop(0) + + self.mediapipe_latencies.append(self.media_pipe_pose.latency) + if len(self.mediapipe_latencies) > 30: + self.mediapipe_latencies.pop(0) - self.set_image(self.plot_latency_graph()) + # If the plot is not being generated, start the process + if self.plot_async_result is None: + self.plot_async_result = self.pool.apply_async( + plot_latency_graph, + (self.ui_latencies, self.mediapipe_latencies, self.frame), + ) + # If the plot is ready, set the image + elif self.plot_async_result.ready(): + image = self.plot_async_result.get() + self.set_image(pygame.image.load(image)) + self.plot_async_result = self.pool.apply_async( + plot_latency_graph, + (self.ui_latencies, self.mediapipe_latencies, self.frame), + ) diff --git a/src/app/main.py b/src/app/main.py index 222ad8a..6c0794c 100644 --- a/src/app/main.py +++ b/src/app/main.py @@ -1,3 +1,5 @@ +from multiprocessing import Pool + import pygame import pygame.camera @@ -36,28 +38,32 @@ def main(): media_pipe_pose=media_pipe_pose, ) - LatencyGraph( - relative_rect=Rect(800, 50, 400, 300), - image_surface=pygame.Surface((400, 300)), - ) - - frame = 0 - clock = pygame.time.Clock() - - while True: - frame += 1 - time_delta = clock.tick(60) / 1000.0 - for event in pygame.event.get(): - if event.type == pygame.QUIT: - return - - manager.process_events(event) - - window_surface.fill(pygame.Color("#000000")) - manager.update(time_delta) - manager.draw_ui(window_surface) - - pygame.display.update() + # Apply multiprocessing using Pool + # Can be used to execute long-running tasks in parallel without blocking the main thread + with Pool(processes=4) as pool: + LatencyGraph( + relative_rect=Rect(800, 50, 400, 300), + image_surface=pygame.Surface((400, 300)), + pool=pool, + media_pipe_pose=media_pipe_pose, + ) + + frame = 0 + clock = pygame.time.Clock() + while True: + frame += 1 + time_delta_ms = clock.tick(60) + for event in pygame.event.get(): + if event.type == pygame.QUIT: + return + + manager.process_events(event) + + window_surface.fill(pygame.Color("#000000")) + manager.update(time_delta_ms / 1000.0) + manager.draw_ui(window_surface) + + pygame.display.update() if __name__ == "__main__": diff --git a/src/app/marker_location_graph.py b/src/app/marker_location_graph.py new file mode 100644 index 0000000..a99c5c7 --- /dev/null +++ b/src/app/marker_location_graph.py @@ -0,0 +1,7 @@ +from pygame_gui.elements import UIImage + + +class MarkerLocationGraph(UIImage): + """ + Class to display the marker location as a graph using 3D coordinates as lines + """