diff --git a/.gitignore b/.gitignore index 5f80f2e..f45ec60 100644 --- a/.gitignore +++ b/.gitignore @@ -166,3 +166,4 @@ cython_debug/ *.jpeg skellytracker/utilities/quine_directory_printer/output/* +recorded_objects.npy diff --git a/skellytracker/RUN_ME.py b/skellytracker/RUN_ME.py index aa18700..cccc56f 100644 --- a/skellytracker/RUN_ME.py +++ b/skellytracker/RUN_ME.py @@ -34,6 +34,7 @@ def main(demo_tracker: str = "mediapipe_holistic_tracker"): charuco_squares_y = 5 number_of_charuco_markers = (charuco_squares_x - 1) * (charuco_squares_y - 1) charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + CharucoTracker( tracked_object_names=charuco_ids, squares_x=charuco_squares_x, diff --git a/skellytracker/SINGLE_IMAGE_RUN.py b/skellytracker/SINGLE_IMAGE_RUN.py index 5ef6a2c..c8d5c57 100644 --- a/skellytracker/SINGLE_IMAGE_RUN.py +++ b/skellytracker/SINGLE_IMAGE_RUN.py @@ -30,6 +30,7 @@ charuco_squares_y = 5 number_of_charuco_markers = (charuco_squares_x - 1) * (charuco_squares_y - 1) charuco_ids = [str(index) for index in range(number_of_charuco_markers)] + CharucoTracker( tracked_object_names=charuco_ids, squares_x=charuco_squares_x, diff --git a/skellytracker/process_folder_of_videos.py b/skellytracker/process_folder_of_videos.py index 4ca1fc8..8e9b924 100644 --- a/skellytracker/process_folder_of_videos.py +++ b/skellytracker/process_folder_of_videos.py @@ -6,7 +6,9 @@ from pydantic import BaseModel +from skellytracker.system.constants import BASE_2D_FILE_NAME from skellytracker.trackers.base_tracker.base_tracker import BaseTracker +from skellytracker.trackers.base_tracker.model_info import ModelInfo from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import ( BrightestPointTracker, ) @@ -37,16 +39,18 @@ logger = logging.getLogger(__name__) -file_name_dictionary = { - "MediapipeHolisticTracker": "mediapipe2dData_numCams_numFrames_numTrackedPoints_pixelXY.npy", - "YOLOMediapipeComboTracker": "mediapipe2dData_numCams_numFrames_numTrackedPoints_pixelXY.npy", - "YOLOPoseTracker": "yolo2dData_numCams_numFrames_numTrackedPoints_pixelXY.npy", - "BrightestPointTracker": "brightestPoint2dData_numCams_numFrames_numTrackedPoints_pixelXY.npy", -} +try: + from skellytracker.trackers.openpose_tracker.openpose_tracker import ( + OpenPoseTracker, + ) +except ModuleNotFoundError: + print("To use openpose_tracker, install skellytracker[openpose]") + +logger = logging.getLogger(__name__) def process_folder_of_videos( - tracker_name: str, + model_info: ModelInfo, tracking_params: BaseModel, synchronized_video_path: Path, output_folder_path: Optional[Path] = None, @@ -57,7 +61,7 @@ def process_folder_of_videos( Process a folder of synchronized videos with the given tracker. Tracked data will be saved to a .npy file with the shape (numCams, numFrames, numTrackedPoints, pixelXYZ). - :param tracker_name: Tracker to use. + :param model_info: Model info for tracker. :param tracking_params: Tracking parameters to use. :param synchronized_video_path: Path to folder of synchronized videos. :param output_folder_path: Path to save tracked data to. @@ -72,7 +76,7 @@ def process_folder_of_videos( else: num_processes = min(num_processes, len(video_paths), cpu_count() - 1) - file_name = file_name_dictionary[tracker_name] + file_name = model_info.name + "_" + BASE_2D_FILE_NAME synchronized_video_path = Path(synchronized_video_path) if output_folder_path is None: output_folder_path = ( @@ -89,7 +93,7 @@ def process_folder_of_videos( annotated_video_path.mkdir(parents=True, exist_ok=True) tasks = [ - (tracker_name, tracking_params, video_path, annotated_video_path) + (model_info.tracker_name, tracking_params, video_path, annotated_video_path) for video_path in video_paths ] @@ -126,9 +130,14 @@ def process_single_video( :param annotated_video_path: Path to save annotated video to. :return: Array of tracking data """ - video_name = ( - video_path.stem + "_mediapipe.mp4" - ) # TODO: fix it so blender output doesn't require mediapipe addendum here + + if tracker_name == "OpenPoseTracker": + video_name = video_path.stem + "_openpose.avi" + else: + video_name = ( + video_path.stem + "_mediapipe.mp4" + ) # TODO: fix it so blender output doesn't require mediapipe addendum here + tracker = get_tracker(tracker_name=tracker_name, tracking_params=tracking_params) logger.info( f"Processing video: {video_name} with tracker: {tracker.__class__.__name__}" @@ -137,7 +146,7 @@ def process_single_video( input_video_filepath=video_path, output_video_filepath=annotated_video_path / video_name, save_data_bool=False, - ) + ) # TODO: raise a custom error here if output_array is None? return output_array @@ -177,6 +186,17 @@ def get_tracker(tracker_name: str, tracking_params: BaseModel) -> BaseTracker: elif tracker_name == "BrightestPointTracker": tracker = BrightestPointTracker() + elif tracker_name == "OpenPoseTracker": + tracker = OpenPoseTracker( + openpose_root_folder_path=tracking_params.openpose_root_folder_path, + output_json_folder_path=tracking_params.output_json_path, + net_resolution=tracking_params.net_resolution, + number_people_max=tracking_params.number_people_max, + track_faces=tracking_params.track_face, + track_hands=tracking_params.track_hands, + output_resolution=tracking_params.output_resolution, + ) + else: raise ValueError("Invalid tracker type") @@ -192,19 +212,26 @@ def get_tracker_params(tracker_name: str) -> BaseModel: return YOLOTrackingParams() elif tracker_name == "BrightestPointTracker": return BaseModel() + elif tracker_name == "OpenPoseTracker": + raise ValueError( + "OpenPoseTracker requires explicitly setting the OpenPose root folder path and output json path, please provide tracking params directly" + ) else: raise ValueError("Invalid tracker type") if __name__ == "__main__": + from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import MediapipeModelInfo + synchronized_video_path = Path( - "/Users/philipqueen/freemocap_data/recording_sessions/freemocap_sample_data/synchronized_videos" + "/Your/Path/To/freemocap_data/recording_sessions/freemocap_sample_data/synchronized_videos" ) + tracker_name = "YOLOMediapipeComboTracker" num_processes = None process_folder_of_videos( - tracker_name=tracker_name, + model_info=MediapipeModelInfo(), tracking_params=get_tracker_params(tracker_name=tracker_name), synchronized_video_path=synchronized_video_path, num_processes=num_processes, diff --git a/skellytracker/system/constants.py b/skellytracker/system/constants.py new file mode 100644 index 0000000..66ead3e --- /dev/null +++ b/skellytracker/system/constants.py @@ -0,0 +1 @@ +BASE_2D_FILE_NAME = "2dData_numCams_numFrames_numTrackedPoints_pixelXY.npy" \ No newline at end of file diff --git a/skellytracker/tests/test_mediapipe_holistic_tracker.py b/skellytracker/tests/test_mediapipe_holistic_tracker.py index 8c53bdb..13c94f3 100644 --- a/skellytracker/tests/test_mediapipe_holistic_tracker.py +++ b/skellytracker/tests/test_mediapipe_holistic_tracker.py @@ -48,7 +48,7 @@ def test_record(test_image): assert processed_results is not None assert processed_results.shape == ( 1, - MediapipeModelInfo.num_tracked_points_total, + MediapipeModelInfo.num_tracked_points, 3, ) diff --git a/skellytracker/tests/test_yolo_mediapipe_combo_tracker.py b/skellytracker/tests/test_yolo_mediapipe_combo_tracker.py index 72bd8fb..f292031 100644 --- a/skellytracker/tests/test_yolo_mediapipe_combo_tracker.py +++ b/skellytracker/tests/test_yolo_mediapipe_combo_tracker.py @@ -58,7 +58,7 @@ def test_record_no_buffer(test_image): assert processed_results is not None assert processed_results.shape == ( 1, - MediapipeModelInfo.num_tracked_points_total, + MediapipeModelInfo.num_tracked_points, 3, ) @@ -122,7 +122,7 @@ def test_record_buffer_by_image_size(test_image): assert processed_results is not None assert processed_results.shape == ( 1, - MediapipeModelInfo.num_tracked_points_total, + MediapipeModelInfo.num_tracked_points, 3, ) @@ -186,7 +186,7 @@ def test_record_buffer_by_box_size(test_image): assert processed_results is not None assert processed_results.shape == ( 1, - MediapipeModelInfo.num_tracked_points_total, + MediapipeModelInfo.num_tracked_points, 3, ) diff --git a/skellytracker/trackers/base_tracker/base_recorder.py b/skellytracker/trackers/base_tracker/base_recorder.py index 9a0fc6e..9cf2c9d 100644 --- a/skellytracker/trackers/base_tracker/base_recorder.py +++ b/skellytracker/trackers/base_tracker/base_recorder.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod import logging -from typing import Dict, Optional +from pathlib import Path +from typing import Dict, Union, Optional import numpy as np @@ -26,7 +27,6 @@ def record( Record the tracked objects as they are created by the tracker. :param tracked_object: A tracked objects dictionary. - :param annotated_image: Image array with tracking results annotated. :return: None """ pass @@ -45,7 +45,7 @@ def clear_recorded_objects(self): self.recorded_objects = [] self.recorded_objects_array = None - def save(self, file_path: str) -> None: + def save(self, file_path: Union[str, Path]) -> None: """ Save the recorded objects to a file. @@ -58,3 +58,19 @@ def save(self, file_path: str) -> None: recorded_objects_array = self.recorded_objects_array logger.info(f"Saving recorded objects to {file_path}") np.save(file_path, recorded_objects_array) + + +class BaseCumulativeRecorder(BaseRecorder): + """ + A base class for recording data from cumulative trackers. + Throws a descriptive error for methods that do not apply to recording data from this type of tracker. + Trackers implementing this will only use the process_tracked_objects method to get data in the proper format. + """ + + def __init__(self): + super().__init__() + + def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: + raise NotImplementedError( + "This tracker does not support by frame recording, please use process_tracked_objects instead" + ) diff --git a/skellytracker/trackers/base_tracker/base_tracker.py b/skellytracker/trackers/base_tracker/base_tracker.py index 4fd9913..b097b36 100644 --- a/skellytracker/trackers/base_tracker/base_tracker.py +++ b/skellytracker/trackers/base_tracker/base_tracker.py @@ -7,7 +7,7 @@ from tqdm import tqdm -from skellytracker.trackers.base_tracker.base_recorder import BaseRecorder +from skellytracker.trackers.base_tracker.base_recorder import BaseCumulativeRecorder, BaseRecorder from skellytracker.trackers.base_tracker.tracked_object import TrackedObject from skellytracker.trackers.base_tracker.video_handler import VideoHandler from skellytracker.trackers.demo_viewers.image_demo_viewer import ImageDemoViewer @@ -65,7 +65,7 @@ def process_video( output_video_filepath: Optional[Union[str, Path]] = None, save_data_bool: bool = False, use_tqdm: bool = True, - ) -> Optional[np.ndarray]: + ) -> Union[np.ndarray, None]: """ Run the tracker on a video. @@ -73,7 +73,7 @@ def process_video( :param output_video_filepath: Path to save annotated video to, does not save video if None. :param save_data_bool: Whether to save the data to a file. :param use_tqdm: Whether to use tqdm to show a progress bar - :return: Array of tracked keypoint data, if save_data_bool is True + :return: Array of tracked keypoint data if tracker has an associated recorder """ cap = cv2.VideoCapture(str(input_video_filepath)) @@ -160,3 +160,56 @@ def image_demo(self, image_path: Path) -> None: image_viewer = ImageDemoViewer(self, self.__class__.__name__) image_viewer.run(image_path=image_path) + + +class BaseCumulativeTracker(BaseTracker): + """ + A base class for tracking algorithms that run cumulatively, i.e are not able to process videos frame by frame. + Throws a descriptive error for the abstract methods of BaseTracker that do not apply to this type of tracker. + Trackers inheriting from this will need to overwrite the `process_video` method. + """ + + def __init__( + self, + tracked_object_names: List[str], + recorder: BaseCumulativeRecorder, + **data: Any, + ): + super().__init__( + tracked_object_names=tracked_object_names, recorder=recorder, **data + ) + + def process_image(self, **kwargs) -> None: + raise NotImplementedError( + "This tracker does not support processing individual images, please use process_video instead." + ) + + def annotate_image(self, **kwargs) -> None: + raise NotImplementedError( + "This tracker does not support processing individual images, please use process_video instead." + ) + + @abstractmethod + def process_video( + self, + input_video_filepath: Union[str, Path], + output_video_filepath: Optional[Union[str, Path]] = None, + save_data_bool: bool = False, + use_tqdm: bool = True, + **kwargs, + ) -> Union[np.ndarray, None]: + """ + Run the tracker on a video. + + :param input_video_filepath: Path to video file. + :param output_video_filepath: Path to save annotated video to, does not save video if None. + :param save_data_bool: Whether to save the data to a file. + :param use_tqdm: Whether to use tqdm to show a progress bar + :return: Array of tracked keypoint data + """ + pass + + def image_demo(self, image_path: Path) -> None: + raise NotImplementedError( + "This tracker does not support processing individual images, please use process_video instead." + ) diff --git a/skellytracker/trackers/base_tracker/base_tracking_params.py b/skellytracker/trackers/base_tracker/base_tracking_params.py index d445647..90b741c 100644 --- a/skellytracker/trackers/base_tracker/base_tracking_params.py +++ b/skellytracker/trackers/base_tracker/base_tracking_params.py @@ -3,4 +3,4 @@ class BaseTrackingParams(BaseModel): num_processes: int = 1 - run_image_tracking: bool = True \ No newline at end of file + run_image_tracking: bool = True diff --git a/skellytracker/trackers/base_tracker/model_info.py b/skellytracker/trackers/base_tracker/model_info.py new file mode 100644 index 0000000..f4e2695 --- /dev/null +++ b/skellytracker/trackers/base_tracker/model_info.py @@ -0,0 +1,13 @@ +from typing import Dict, List, Optional, Union + + +class ModelInfo(dict): + name: str + tracker_name: str + landmark_names: List[str] + num_tracked_points: int + tracked_object_names: Optional[list] = None + virtual_markers_definitions: Optional[Dict[str, Dict[str, List[Union[str, float]]]]] = None + segment_connections: Optional[Dict[str, Dict[str, str]]] = None + center_of_mass_definitions: Optional[Dict[str, Dict[str, float]]] = None + joint_hierarchy: Optional[Dict[str, List[str]]] = None diff --git a/skellytracker/trackers/base_tracker/video_handler.py b/skellytracker/trackers/base_tracker/video_handler.py index 9a2087a..13d7249 100644 --- a/skellytracker/trackers/base_tracker/video_handler.py +++ b/skellytracker/trackers/base_tracker/video_handler.py @@ -25,7 +25,9 @@ def __init__( """ self.output_path = output_path fourcc = cv2.VideoWriter.fourcc(*codec) - self.video_writer = cv2.VideoWriter(str(output_path), fourcc, fps, frame_size) + self.video_writer = cv2.VideoWriter( + str(output_path), fourcc, fps, frame_size + ) def add_frame(self, frame: np.ndarray) -> None: """ diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py index f6254b2..0a48857 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_recorder.py @@ -14,7 +14,7 @@ def record(self, tracked_objects: Dict[str, TrackedObject]) -> None: self.recorded_objects.append( [ deepcopy(tracked_objects[tracked_object_name]) - for tracked_object_name in MediapipeModelInfo.mediapipe_tracked_object_names + for tracked_object_name in MediapipeModelInfo.tracked_object_names ] ) @@ -27,7 +27,7 @@ def process_tracked_objects(self, **kwargs) -> np.ndarray: self.recorded_objects_array = np.zeros( ( len(self.recorded_objects), - MediapipeModelInfo.num_tracked_points_total, + MediapipeModelInfo.num_tracked_points, 3, ) ) diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py index 84c34c6..5c0ad66 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_holistic_tracker.py @@ -23,7 +23,7 @@ def __init__( smooth_landmarks=True, ): super().__init__( - tracked_object_names=MediapipeModelInfo.mediapipe_tracked_object_names, + tracked_object_names=MediapipeModelInfo.tracked_object_names, recorder=MediapipeHolisticRecorder(), ) self.mp_drawing = mp.solutions.drawing_utils diff --git a/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py b/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py index 2dc0320..5e28369 100644 --- a/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py +++ b/skellytracker/trackers/mediapipe_tracker/mediapipe_model_info.py @@ -1,33 +1,189 @@ -from typing import Literal +from typing import List, Literal from mediapipe.python.solutions import holistic as mp_holistic from mediapipe.python.solutions.face_mesh import FACEMESH_NUM_LANDMARKS_WITH_IRISES from skellytracker.trackers.base_tracker.base_tracking_params import BaseTrackingParams +from skellytracker.trackers.base_tracker.model_info import ModelInfo -mediapipe_body_landmark_names = [ - landmark.name.lower() for landmark in mp_holistic.PoseLandmark -] -mediapipe_hand_landmark_names = [ - landmark.name.lower() for landmark in mp_holistic.HandLandmark -] - -class MediapipeModelInfo: - num_tracked_points_body = len(mediapipe_body_landmark_names) +# values for segment weight and segment mass percentages taken from Winter anthropometry tables +# https://imgur.com/a/aD74j +# Winter, D.A. (2005) Biomechanics and Motor Control of Human Movement. 3rd Edition, John Wiley & Sons, Inc., Hoboken. +class MediapipeModelInfo(ModelInfo): + name = "mediapipe" + tracker_name = "MediapipeHolisticTracker" + body_landmark_names = [ + landmark.name.lower() for landmark in mp_holistic.PoseLandmark + ] + hand_landmark_names = [ + landmark.name.lower() for landmark in mp_holistic.HandLandmark + ] + face_landmark_names = [ + "right_eye", + "left_eye", + "nose_tip", + "mouth_center", + "right_ear_tragion", + "left_ear_tragion", + ] + landmark_names = body_landmark_names + hand_landmark_names + face_landmark_names + num_tracked_points_body = len(body_landmark_names) num_tracked_points_face = FACEMESH_NUM_LANDMARKS_WITH_IRISES - num_tracked_points_left_hand = len(mediapipe_hand_landmark_names) - num_tracked_points_right_hand = len(mediapipe_hand_landmark_names) - num_tracked_points_total = ( - len(mediapipe_body_landmark_names) - + 2 * len(mediapipe_hand_landmark_names) - + FACEMESH_NUM_LANDMARKS_WITH_IRISES + num_tracked_points_left_hand = len(hand_landmark_names) + num_tracked_points_right_hand = len(hand_landmark_names) + num_tracked_points = ( + len(body_landmark_names) + + 2 * len(hand_landmark_names) + + num_tracked_points_face ) - mediapipe_tracked_object_names = [ + tracked_object_names = [ "pose_landmarks", "right_hand_landmarks", "left_hand_landmarks", "face_landmarks", ] + virtual_markers_definitions = { + "head_center": { + "marker_names": ["left_ear", "right_ear"], + "marker_weights": [0.5, 0.5], + }, + "neck_center": { + "marker_names": ["left_shoulder", "right_shoulder"], + "marker_weights": [0.5, 0.5], + }, + "trunk_center": { + "marker_names": [ + "left_shoulder", + "right_shoulder", + "left_hip", + "right_hip", + ], + "marker_weights": [0.25, 0.25, 0.25, 0.25], + }, + "hips_center": { + "marker_names": ["left_hip", "right_hip"], + "marker_weights": [0.5, 0.5], + }, + } + segment_connections = { + "head": {"proximal": "left_ear", "distal": "right_ear"}, + "neck": { + "proximal": "head_center", + "distal": "neck_center", + }, + "spine": { + "proximal": "neck_center", + "distal": "hips_center", + }, + "right_shoulder": {"proximal": "neck_center", "distal": "right_shoulder"}, + "left_shoulder": {"proximal": "neck_center", "distal": "left_shoulder"}, + "right_upper_arm": {"proximal": "right_shoulder", "distal": "right_elbow"}, + "left_upper_arm": {"proximal": "left_shoulder", "distal": "left_elbow"}, + "right_forearm": {"proximal": "right_elbow", "distal": "right_wrist"}, + "left_forearm": {"proximal": "left_elbow", "distal": "left_wrist"}, + "right_hand": {"proximal": "right_wrist", "distal": "right_index"}, + "left_hand": {"proximal": "left_wrist", "distal": "left_index"}, + "right_pelvis": {"proximal": "hips_center", "distal": "right_hip"}, + "left_pelvis": {"proximal": "hips_center", "distal": "left_hip"}, + "right_thigh": {"proximal": "right_hip", "distal": "right_knee"}, + "left_thigh": {"proximal": "left_hip", "distal": "left_knee"}, + "right_shank": {"proximal": "right_knee", "distal": "right_ankle"}, + "left_shank": {"proximal": "left_knee", "distal": "left_ankle"}, + "right_foot": {"proximal": "right_ankle", "distal": "right_foot_index"}, + "left_foot": {"proximal": "left_ankle", "distal": "left_foot_index"}, + "right_heel": {"proximal": "right_ankle", "distal": "right_heel"}, + "left_heel": {"proximal": "left_ankle", "distal": "left_heel"}, + "right_foot_bottom": {"proximal": "right_heel", "distal": "right_foot_index"}, + "left_foot_bottom": {"proximal": "left_heel", "distal": "left_foot_index"}, + } + center_of_mass_definitions = { + "head": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.081, + }, + "spine": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.497, + }, + "right_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "left_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "right_forearm": { + "segment_com_length": 0.430, + "segment_com_percentage": 0.016, + }, + "left_forearm": { + "segment_com_length": 0.430, + "segment_com_percentage": 0.016, + }, + "right_hand": { + "segment_com_length": 0.506, + "segment_com_percentage": 0.006, + }, + "left_hand": { + "segment_com_length": 0.506, + "segment_com_percentage": 0.006, + }, + "right_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "left_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "right_shank": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.0465, + }, + "left_shank": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.0465, + }, + "right_foot": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.0145, + }, + "left_foot": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.0145, + }, + } + joint_hierarchy = { + "hips_center": ["left_hip", "right_hip", "trunk_center"], + "trunk_center": ["neck_center"], + "neck_center": ["left_shoulder", "right_shoulder", "head_center"], + "head_center": [ + "nose", + "left_eye_inner", + "left_eye", + "left_eye_outer", + "right_eye_inner", + "right_eye", + "right_eye_outer", + "left_ear", + "right_ear", + "mouth_left", + "mouth_right", + ], + "left_shoulder": ["left_elbow"], + "left_elbow": ["left_wrist"], + "left_wrist": ["left_pinky", "left_index", "left_thumb"], + "right_shoulder": ["right_elbow"], + "right_elbow": ["right_wrist"], + "right_wrist": ["right_pinky", "right_index", "right_thumb"], + "left_hip": ["left_knee"], + "left_knee": ["left_ankle"], + "left_ankle": ["left_heel", "left_foot_index"], + "right_hip": ["right_knee"], + "right_knee": ["right_ankle"], + "right_ankle": ["right_heel", "right_foot_index"], + } class MediapipeTrackingParams(BaseTrackingParams): @@ -35,7 +191,7 @@ class MediapipeTrackingParams(BaseTrackingParams): mediapipe_model_complexity: int = 2 min_detection_confidence: float = 0.5 min_tracking_confidence: float = 0.5 - static_image_mode: bool = False + static_image_mode: bool = True yolo_model_size: Literal[ "nano", "small", "medium", "large", "extra_large", "high_res" ] = "nano" @@ -43,3 +199,53 @@ class MediapipeTrackingParams(BaseTrackingParams): buffer_size_method: Literal["buffer_by_box_size", "buffer_by_image_size"] = ( "buffer_by_box_size" ) + + +def mediapipe_body_names_match_expected( + mediapipe_body_landmark_names: List[str], +) -> bool: + """ + Check if the mediapipe folks have changed their landmark names. If they have, then this function may need to be updated. + + Args: + mediapipe_body_landmark_names: List of strings, each string is the name of a mediapipe landmark. + + Returns: + bool: True if the mediapipe landmark names are as expected, False otherwise. + """ + expected_mediapipe_body_landmark_names = [ + "nose", + "left_eye_inner", + "left_eye", + "left_eye_outer", + "right_eye_inner", + "right_eye", + "right_eye_outer", + "left_ear", + "right_ear", + "mouth_left", + "mouth_right", + "left_shoulder", + "right_shoulder", + "left_elbow", + "right_elbow", + "left_wrist", + "right_wrist", + "left_pinky", + "right_pinky", + "left_index", + "right_index", + "left_thumb", + "right_thumb", + "left_hip", + "right_hip", + "left_knee", + "right_knee", + "left_ankle", + "right_ankle", + "left_heel", + "right_heel", + "left_foot_index", + "right_foot_index", + ] + return mediapipe_body_landmark_names == expected_mediapipe_body_landmark_names diff --git a/skellytracker/trackers/openpose_tracker/openpose_model_info.py b/skellytracker/trackers/openpose_tracker/openpose_model_info.py new file mode 100644 index 0000000..8befd3e --- /dev/null +++ b/skellytracker/trackers/openpose_tracker/openpose_model_info.py @@ -0,0 +1,164 @@ +from skellytracker.trackers.base_tracker.base_tracking_params import BaseTrackingParams +from skellytracker.trackers.base_tracker.model_info import ModelInfo + +from typing import Optional + + +class OpenPoseModelInfo(ModelInfo): + name = "openpose" + tracker_name = "OpenPoseTracker" + body_landmark_names = [ + "nose", + "neck", + "right_shoulder", + "right_elbow", + "right_wrist", + "left_shoulder", + "left_elbow", + "left_wrist", + "hip_center", + "right_hip", + "right_knee", + "right_ankle", + "left_hip", + "left_knee", + "left_ankle", + "right_eye", + "left_eye", + "right_ear", + "left_ear", + "left_big_toe", + "left_small_toe", + "left_heel", + "right_big_toe", + "right_small_toe", + "right_heel" + ] + landmark_names = body_landmark_names + num_tracked_points_body = len(body_landmark_names) + num_tracked_points_face = 70 + num_tracked_points_left_hand = 21 + num_tracked_points_right_hand = 21 + + num_tracked_points = ( + num_tracked_points_body + + num_tracked_points_left_hand + + num_tracked_points_right_hand + + num_tracked_points_face + ) + tracked_object_names = ["pose_landmarks"] + virtual_markers_definitions = { + "head_center": { + "marker_names": ["left_ear", "right_ear"], + "marker_weights": [0.5, 0.5], + }, + "trunk_center": { + "marker_names": [ + "left_shoulder", + "right_shoulder", + "left_hip", + "right_hip", + ], + "marker_weights": [0.25, 0.25, 0.25, 0.25], + }, + } + segment_connections = { + "head": {"proximal": "left_ear", "distal": "right_ear"}, + "neck": {"proximal": "head_center", "distal": "neck"}, + "spine": {"proximal": "neck", "distal": "hip_center"}, + "right_shoulder": {"proximal": "neck", "distal": "right_shoulder"}, + "left_shoulder": {"proximal": "neck", "distal": "left_shoulder"}, + "right_upper_arm": {"proximal": "right_shoulder", "distal": "right_elbow"}, + "left_upper_arm": {"proximal": "left_shoulder", "distal": "left_elbow"}, + "right_forearm": {"proximal": "right_elbow", "distal": "right_wrist"}, + "left_forearm": {"proximal": "left_elbow", "distal": "left_wrist"}, + "right_pelvis": {"proximal": "hip_center", "distal": "right_hip"}, + "left_pelvis": {"proximal": "hip_center", "distal": "left_hip"}, + "right_thigh": {"proximal": "right_hip", "distal": "right_knee"}, + "left_thigh": {"proximal": "left_hip", "distal": "left_knee"}, + "right_shank": {"proximal": "right_knee", "distal": "right_ankle"}, + "left_shank": {"proximal": "left_knee", "distal": "left_ankle"}, + "right_foot": {"proximal": "right_ankle", "distal": "right_big_toe"}, + "left_foot": {"proximal": "left_ankle", "distal": "left_big_toe"}, + "right_heel": {"proximal": "right_ankle", "distal": "right_heel"}, + "left_heel": {"proximal": "left_ankle", "distal": "left_heel"}, + "right_foot_bottom": {"proximal": "right_heel", "distal": "right_big_toe"}, + "left_foot_bottom": {"proximal": "left_heel", "distal": "left_big_toe"}, + } + center_of_mass_definitions = { #NOTE: using forearm/hand definition from Winter tables, as we don't have hand definitions here + "head": { + "segment_com_length": .5, + "segment_com_percentage": .081, + }, + "spine": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.497, + }, + "right_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "left_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "right_forearm": { + "segment_com_length": 0.682, + "segment_com_percentage": 0.022, + }, + "left_forearm": { + "segment_com_length": 0.682, + "segment_com_percentage": 0.022, + }, + "right_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "left_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "right_shank": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.0465, + }, + "left_shank": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.0465, + }, + "right_foot": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.0145, + }, + "left_foot": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.0145, + }, + } + joint_hierarchy = { + "hip_center": ["left_hip", "right_hip", "trunk_center"], + "trunk_center": ["neck"], + "neck": ["left_shoulder", "right_shoulder", "head_center"], + "head_center": ["nose", "left_ear", "right_ear", "left_eye", "right_eye"], + "left_shoulder": ["left_elbow"], + "left_elbow": ["left_wrist"], + "right_shoulder": ["right_elbow"], + "right_elbow": ["right_wrist"], + "left_hip": ["left_knee"], + "left_knee": ["left_ankle"], + "left_ankle": ["left_big_toe", "left_small_toe", "left_heel"], + "right_hip": ["right_knee"], + "right_knee": ["right_ankle"], + "right_ankle": ["right_big_toe", "right_small_toe", "right_heel"], + } + + +class OpenPoseTrackingParams(BaseTrackingParams): + openpose_root_folder_path: str + output_json_path: Optional[str] = None + net_resolution: str = "-1x320" + number_people_max: int = 1 + track_hands: bool = True + track_face: bool = True + write_video: bool = True + output_resolution: str = "-1x-1" diff --git a/skellytracker/trackers/openpose_tracker/openpose_recorder.py b/skellytracker/trackers/openpose_tracker/openpose_recorder.py new file mode 100644 index 0000000..4b57d95 --- /dev/null +++ b/skellytracker/trackers/openpose_tracker/openpose_recorder.py @@ -0,0 +1,116 @@ +import json +from typing import Dict, Union +import numpy as np +from pathlib import Path +import re +from tqdm import tqdm + +from skellytracker.trackers.base_tracker.base_recorder import BaseCumulativeRecorder +from skellytracker.trackers.openpose_tracker.openpose_model_info import ( + OpenPoseModelInfo, +) + + +class OpenPoseRecorder(BaseCumulativeRecorder): + def __init__( + self, + track_hands: bool = False, + track_faces: bool = False, + ): + super().__init__() + self.track_hands = track_hands + self.track_faces = track_faces + + def extract_frame_index(self, filename: str) -> Union[int, None]: + """Extract the numeric part indicating the frame index from the filename.""" + match = re.search(r"_(\d{12})_keypoints", filename) + return int(match.group(1)) if match else None + + def parse_openpose_jsons(self, json_directory: Union[Path, str]) -> np.ndarray: + # Remove the iteration over subdirectories and focus on a single directory + json_directory = Path(json_directory) + files = list(Path(json_directory).glob("*.json")) + num_frames = len(files) + frame_indices = [ + index + for f in files + if (index := self.extract_frame_index(f.name)) is not None + ] + frame_indices.sort() + + if len(frame_indices) != num_frames: + raise ValueError( + f"Invalid number of frames in {json_directory}: expected {num_frames} != {len(frame_indices)} frames in file" + ) + + num_markers = OpenPoseModelInfo.num_tracked_points_body + if self.track_hands: + num_markers += ( + OpenPoseModelInfo.num_tracked_points_right_hand + + OpenPoseModelInfo.num_tracked_points_left_hand + ) + if self.track_faces: + num_markers += OpenPoseModelInfo.num_tracked_points_face + + # Initialize a single camera array since we're only processing one video at a time + data_array = np.full((num_frames, num_markers, 3), np.nan) + + # Process each JSON file in the directory + for file_index, json_file in enumerate( + tqdm(files, desc=f"Processing {json_directory.name} JSONs") + ): + with open(json_file) as f: + data = json.load(f) + + if data["people"]: + keypoints = self.extract_keypoints(data["people"][0]) + data_array[frame_indices[file_index], :, :] = keypoints + + return data_array + + def extract_keypoints(self, person_data: Dict[str, np.ndarray]) -> np.ndarray: + """Extract and organize keypoints from person data.""" + + body_markers = OpenPoseModelInfo.num_tracked_points_body + hand_markers = ( + OpenPoseModelInfo.num_tracked_points_left_hand + + OpenPoseModelInfo.num_tracked_points_right_hand + ) + face_markers = OpenPoseModelInfo.num_tracked_points_face + + # Initialize a full array of NaNs for keypoints + keypoints_array = np.full( + (body_markers + (hand_markers) + face_markers, 3), np.nan + ) + + # Populate the array with available data + if "pose_keypoints_2d" in person_data: + keypoints_array[:body_markers, :] = np.reshape( + person_data["pose_keypoints_2d"], (-1, 3) + )[:body_markers, :] + if ( + "hand_left_keypoints_2d" in person_data + and "hand_right_keypoints_2d" in person_data + ): + keypoints_array[body_markers : body_markers + OpenPoseModelInfo.num_tracked_points_left_hand, :] = np.reshape( + person_data["hand_left_keypoints_2d"], (-1, 3) + )[:hand_markers, :] + keypoints_array[ + body_markers + OpenPoseModelInfo.num_tracked_points_left_hand : body_markers + OpenPoseModelInfo.num_tracked_points_left_hand + OpenPoseModelInfo.num_tracked_points_right_hand, : + ] = np.reshape(person_data["hand_right_keypoints_2d"], (-1, 3))[ + :hand_markers, : + ] + if "face_keypoints_2d" in person_data: + keypoints_array[body_markers + hand_markers :, :] = np.reshape( + person_data["face_keypoints_2d"], (-1, 3) + )[:face_markers, :] + + return keypoints_array + + def process_tracked_objects(self, output_json_path: Path) -> np.ndarray: + """ + Convert the recorded JSON data into the structured numpy array format. + """ + # In this case, the recorded_objects are already in the desired format, so we simply return them. + self.recorded_objects_array = self.parse_openpose_jsons(output_json_path) + return self.recorded_objects_array diff --git a/skellytracker/trackers/openpose_tracker/openpose_tracker.py b/skellytracker/trackers/openpose_tracker/openpose_tracker.py new file mode 100644 index 0000000..8a79df4 --- /dev/null +++ b/skellytracker/trackers/openpose_tracker/openpose_tracker.py @@ -0,0 +1,149 @@ +import subprocess +from pathlib import Path +from typing import Optional, Union +from skellytracker.trackers.base_tracker.base_tracker import BaseCumulativeTracker +from skellytracker.trackers.openpose_tracker.openpose_recorder import OpenPoseRecorder + + +class OpenPoseTracker(BaseCumulativeTracker): + def __init__( + self, + openpose_root_folder_path: Union[str, Path], + output_json_folder_path: Optional[Union[str, Path]] = None, + net_resolution: str = "-1x320", + number_people_max: int = 1, + track_hands: bool = True, + track_faces: bool = True, + output_resolution: str = "-1x-1", + ): + """ + Initialize the OpenPoseTracker. + + :param recorder: An instance of OpenPoseRecorder for handling the output. + :param openpose_root_folder_path: Path to the OpenPose root folder. + :param output_json_folder_path: Path to the output JSON folder. + :param net_resolution: Network resolution for OpenPose processing. + :param number_people_max: Maximum number of people to detect. + :param track_hands: Whether to track hands. + :param track_faces: Whether to track faces. + :param output_resolution: Output resolution for video. + """ + super().__init__( + tracked_object_names=[], + recorder=OpenPoseRecorder( + track_hands=track_hands, + track_faces=track_faces, + ), + track_hands=track_hands, + track_faces=track_faces, + ) + self.openpose_root_folder_path = Path(openpose_root_folder_path) + self.output_json_folder_path = output_json_folder_path + self.net_resolution = net_resolution + self.number_people_max = number_people_max + self.track_hands = track_hands + self.track_faces = track_faces + self.output_resolution = output_resolution + + def set_track_hands(self, track_hands: bool): + self._track_hands = track_hands + self.recorder.track_hands = track_hands + + def set_track_faces(self, track_faces: bool): + self._track_faces = track_faces + self.recorder.track_faces = track_faces + + def set_json_output_path(self, output_json_folder_path: Union[str, Path]): + self.output_json_folder_path = Path(output_json_folder_path) + + def process_video( + self, + input_video_filepath: Union[str, Path], + output_video_filepath: Union[str, Path], + save_data_bool: bool = False, + use_tqdm: bool = True, # TODO: this is unused, replace with an openpose flag or remove + **kwargs, + ): + """ + Run the OpenPose demo on a video file to generate JSON outputs + in a unique directory for each video. + + :param input_video_filepath: Path to the input video file. + :param output_video_filepath: Path to the output video file. + :param save_data_bool: Whether to save the data. + :param use_tqdm: Whether to use tqdm progress bar. + :return: The output array, or None if recorder isn't initialized in tracker. + """ + # Extract video name without extension to use as a unique folder name + video_name = Path(input_video_filepath).stem + + if self.output_json_folder_path is None: + self.output_json_folder_path = Path(input_video_filepath).parent.parent / "output_data" / "raw_data" / "openpose_jsons" + + Path(self.output_json_folder_path).mkdir(parents=True, exist_ok=True) + + unique_json_output_path = Path(self.output_json_folder_path) / video_name + unique_json_output_path.mkdir(parents=True, exist_ok=True) + + # Full path to the OpenPose executable + openpose_executable_path = ( + self.openpose_root_folder_path / "bin" / "OpenPoseDemo.exe" + ) + + openpose_command = [ + str(openpose_executable_path), # Full path to the OpenPose executable + "--video", + str(input_video_filepath), + "--write_json", + str(unique_json_output_path), + "--net_resolution", + str(self.net_resolution), + "--number_people_max", + str(self.number_people_max), + "--write_video", + str(output_video_filepath), + "--output_resolution", + str(self.output_resolution), + ] + + if self.track_hands: + openpose_command.append("--hand") + if self.track_faces: + openpose_command.append("--face") + + # Update the subprocess command to use the unique output directory + try: + subprocess.run( # noqa: S603 + openpose_command, + shell=False, + cwd=self.openpose_root_folder_path, # Set the current working directory for the subprocess + check=True, + ) + except subprocess.CalledProcessError as e: + print(f"Error: {e}") + return None + + if self.recorder is not None: + output_array = self.recorder.process_tracked_objects( + output_json_path=unique_json_output_path + ) + if save_data_bool: + self.recorder.save( + file_path=str(Path(input_video_filepath).with_suffix(".npy")) + ) + else: + output_array = None + + return output_array + + +if __name__ == "__main__": + # Example usage + openpose_root_folder_path = r"C:\openpose" + input_video_filepath = r'C:\path\to\input\video.mp4' + output_video_filepath = r'C:\path\to\output\video.avi' + + tracker = OpenPoseTracker( + openpose_root_folder_path=str(openpose_root_folder_path), + ) + tracker.process_video(input_video_filepath, output_video_filepath) diff --git a/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py b/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py index e0cb3cb..b0cd08b 100644 --- a/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py +++ b/skellytracker/trackers/yolo_mediapipe_combo_tracker/yolo_mediapipe_combo_tracker.py @@ -34,7 +34,7 @@ def __init__( ] = "buffer_by_box_size", ): super().__init__( - tracked_object_names=MediapipeModelInfo.mediapipe_tracked_object_names, + tracked_object_names=MediapipeModelInfo.tracked_object_names, recorder=MediapipeHolisticRecorder(), ) self.mp_drawing = mp.solutions.drawing_utils diff --git a/skellytracker/trackers/yolo_tracker/yolo_model_info.py b/skellytracker/trackers/yolo_tracker/yolo_model_info.py index 6127299..f0de791 100644 --- a/skellytracker/trackers/yolo_tracker/yolo_model_info.py +++ b/skellytracker/trackers/yolo_tracker/yolo_model_info.py @@ -1,7 +1,10 @@ from skellytracker.trackers.base_tracker.base_tracking_params import BaseTrackingParams +from skellytracker.trackers.base_tracker.model_info import ModelInfo -class YOLOModelInfo: +class YOLOModelInfo(ModelInfo): + name = "yolo" + tracker_name = "YOLOPoseTracker" num_tracked_points = 17 model_dictionary = { "nano": "yolov8n-pose.pt", @@ -10,75 +13,135 @@ class YOLOModelInfo: "large": "yolov8l-pose.pt", "extra_large": "yolov8x-pose.pt", "high_res": "yolov8x-pose-p6.pt", + } # TODO: rename to tracker_dictionary to avoid pydantic 2 conflict? + landmark_names = [ + "nose", + "left_eye", + "right_eye", + "left_ear", + "right_ear", + "left_shoulder", + "right_shoulder", + "left_elbow", + "right_elbow", + "left_wrist", + "right_wrist", + "left_hip", + "right_hip", + "left_knee", + "right_knee", + "left_ankle", + "right_ankle", + ] + virtual_markers_definitions = { + "head_center": { + "marker_names": ["left_ear", "right_ear"], + "marker_weights": [0.5, 0.5], + }, + "neck_center": { + "marker_names": ["left_shoulder", "right_shoulder"], + "marker_weights": [0.5, 0.5], + }, + "trunk_center": { + "marker_names": [ + "left_shoulder", + "right_shoulder", + "left_hip", + "right_hip", + ], + "marker_weights": [0.25, 0.25, 0.25, 0.25], + }, + "hips_center": { + "marker_names": ["left_hip", "right_hip"], + "marker_weights": [0.5, 0.5], + }, } - marker_dict = { - 0: "nose", - 1: "left_eye", - 2: "right_eye", - 3: "left_ear", - 4: "right_ear", - 5: "left_shoulder", - 6: "right_shoulder", - 7: "left_elbow", - 8: "right_elbow", - 9: "left_wrist", - 10: "right_wrist", - 11: "left_hip", - 12: "right_hip", - 13: "left_knee", - 14: "right_knee", - 15: "left_ankle", - 16: "right_ankle", + segment_connections = { + "head": {"proximal": "left_ear", "distal": "right_ear"}, + "neck": { + "proximal": "head_center", + "distal": "neck_center", + }, + "spine": { + "proximal": "neck_center", + "distal": "hips_center", + }, + "right_shoulder": {"proximal": "neck_center", "distal": "right_shoulder"}, + "left_shoulder": {"proximal": "neck_center", "distal": "left_shoulder"}, + "right_upper_arm": {"proximal": "right_shoulder", "distal": "right_elbow"}, + "left_upper_arm": {"proximal": "left_shoulder", "distal": "left_elbow"}, + "right_forearm": {"proximal": "right_elbow", "distal": "right_wrist"}, + "left_forearm": {"proximal": "left_elbow", "distal": "left_wrist"}, + "right_pelvis": {"proximal": "hips_center", "distal": "right_hip"}, + "left_pelvis": {"proximal": "hips_center", "distal": "left_hip"}, + "right_thigh": {"proximal": "right_hip", "distal": "right_knee"}, + "left_thigh": {"proximal": "left_hip", "distal": "left_knee"}, + "right_shank": {"proximal": "right_knee", "distal": "right_ankle"}, + "left_shank": {"proximal": "left_knee", "distal": "left_ankle"}, + } + center_of_mass_definitions = { + "head": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.081, + }, + "spine": { + "segment_com_length": 0.5, + "segment_com_percentage": 0.497, + }, + "right_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "left_upper_arm": { + "segment_com_length": 0.436, + "segment_com_percentage": 0.028, + }, + "right_forearm": { + "segment_com_length": 0.682, + "segment_com_percentage": 0.022, + }, + "left_forearm": { + "segment_com_length": 0.682, + "segment_com_percentage": 0.022, + }, + "right_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "left_thigh": { + "segment_com_length": 0.433, + "segment_com_percentage": 0.1, + }, + "right_shank": { + "segment_com_length": 0.606, + "segment_com_percentage": 0.061, + }, + "left_shank": { + "segment_com_length": 0.606, + "segment_com_percentage": 0.061, + }, + } + joint_hierarchy = { + "hips_center": ["left_hip", "right_hip", "trunk_center"], + "trunk_center": ["neck_center"], + "neck_center": ["left_shoulder", "right_shoulder", "head_center"], + "head_center": [ + "nose", + "left_eye", + "right_eye", + "left_ear", + "right_ear", + ], + "left_shoulder": ["left_elbow"], + "left_elbow": ["left_wrist"], + "right_shoulder": ["right_elbow"], + "right_elbow": ["right_wrist"], + "left_hip": ["left_knee"], + "left_knee": ["left_ankle"], + "right_hip": ["right_knee"], + "right_knee": ["right_ankle"], } - body_segment_names = [ - "head", - "trunk", - "right_upper_arm", - "left_upper_arm", - "right_forearm", - "left_forearm", - "right_thigh", - "left_thigh", - "right_shin", - "left_shin", - ] - joint_connections = [ - ["left_ear", "right_ear"], - ["mid_chest_marker", "mid_hip_marker"], - ["right_shoulder", "right_elbow"], - ["left_shoulder", "left_elbow"], - ["right_elbow", "right_wrist"], - ["left_elbow", "left_wrist"], - ["right_hip", "right_knee"], - ["left_hip", "left_knee"], - ["right_knee", "right_ankle"], - ["left_knee", "left_ankle"], - ] - segment_COM_lengths = [ - 0.5, - 0.5, - 0.436, - 0.436, - 0.682, - 0.682, - 0.433, - 0.433, - 0.606, - 0.606, - ] - segment_COM_percentages = [ - 0.081, - 0.497, - 0.028, - 0.028, - 0.022, - 0.022, - 0.1, - 0.1, - 0.061, - 0.061, - ] class YOLOTrackingParams(BaseTrackingParams): - model_size: str = "medium" + model_size: str = "medium" # TODO: rename to tracker_model_size to avoid pydantic 2 conflict?