From ebfc8bc4c08378eaa91f6412c88ca5938a31ee9f Mon Sep 17 00:00:00 2001 From: MitchellJC <81349046+MitchellJC@users.noreply.github.com> Date: Wed, 14 Aug 2024 11:08:42 +1000 Subject: [PATCH] ref: added typehints for DebugPostureTracker init --- client/models/pose_detection/routines.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/client/models/pose_detection/routines.py b/client/models/pose_detection/routines.py index b227d47..246ec08 100644 --- a/client/models/pose_detection/routines.py +++ b/client/models/pose_detection/routines.py @@ -1,11 +1,17 @@ """Routines that can be integrated into a main control flow.""" from importlib import resources +from typing import Callable, Mapping import cv2 import mediapipe as mp +from mediapipe.framework import calculator_pb2 +from mediapipe.python._framework_bindings.packet import Packet from mediapipe.tasks.python.core.base_options import BaseOptions from mediapipe.tasks.python.vision import RunningMode +from mediapipe.tasks.python.vision.core.vision_task_running_mode import ( + VisionTaskRunningMode, +) from mediapipe.tasks.python.vision.pose_landmarker import ( PoseLandmarker, PoseLandmarkerOptions, @@ -21,12 +27,17 @@ class DebugPostureTracker(PoseLandmarker): """Handles routines for a Debugging Posture Tracker""" - def __init__(self, graph_config, running_mode, packet_callback) -> None: + def __init__( + self, + graph_config: calculator_pb2.CalculatorGraphConfig, + running_mode: VisionTaskRunningMode, + packet_callback: Callable[[Mapping[str, Packet]], None], + ) -> None: super().__init__(graph_config, running_mode, packet_callback) self.annotated_image = AnnotatedImage() self.video_capture = cv2.VideoCapture(0) - def track_posture(self): + def track_posture(self) -> None: """Get frame from video capture device and process with pose model, then posture algorithm. Print debugging info and display landmark annotated frame. """