Aaron/openpose tracking

skellytracker/process_folder_of_videos.py

@@ -13,6 +13,7 @@
 from skellytracker.trackers.bright_point_tracker.brightest_point_tracker import (
     BrightestPointTracker,
 )
+
 from skellytracker.utilities.get_video_paths import get_video_paths
 
 try:
@@ -35,8 +36,24 @@
     from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import (
         MediapipeTrackingParams,
     )
+    from skellytracker.trackers.mediapipe_tracker.mediapipe_model_info import (
+        MediapipeModelInfo,
+    )
 except:
     print("To use mediapipe_holistic_tracker, install skellytracker[mediapipe]")
+
+try:
+    from skellytracker.trackers.openpose_tracker.openpose_tracker import (
+          OpenPoseTracker,
+    )
+    from skellytracker.trackers.openpose_tracker.openpose_model_info import (
+            OpenPoseTrackingParams,
+    )
+    from skellytracker.trackers.openpose_tracker.openpose_model_info import (
+            OpenPoseModelInfo,
+    )
+except:
+    print("To use openpose_tracker, install skellytracker[openpose]")
 
 logger = logging.getLogger(__name__)
 
@@ -122,9 +139,16 @@ 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__}"
@@ -173,6 +197,14 @@ def get_tracker(tracker_name: str, tracking_params: BaseModel) -> BaseTracker:
     elif tracker_name == "BrightestPointTracker":
         tracker = BrightestPointTracker()
 
+    elif tracker_name == 'OpenPoseTracker':
+        tracker = OpenPoseTracker(
+            openpose_exe_path=tracking_params.openpose_exe_path,
+            output_json_path=tracking_params.output_json_path,
+            net_resolution=tracking_params.net_resolution,
+            number_people_max=tracking_params.number_people_max,
+        )
+
     else:
         raise ValueError("Invalid tracker type")
 
@@ -188,20 +220,53 @@ def get_tracker_params(tracker_name: str) -> BaseModel:
         return YOLOTrackingParams()
     elif tracker_name == "BrightestPointTracker":
         return BaseModel()
+    elif tracker_name == 'OpenPoseTracker':
+        return OpenPoseTrackingParams()
     else:
         raise ValueError("Invalid tracker type")
 
 
 if __name__ == "__main__":
-    synchronized_video_path = Path(
-        "/Users/philipqueen/freemocap_data/recording_sessions/freemocap_sample_data/synchronized_videos"
-    )
-    tracker_name = "YOLOMediapipeComboTracker"
-    num_processes = None
+
+    # synchronized_video_path = Path(
+    #     "/Users/philipqueen/freemocap_data/recording_sessions/freemocap_sample_data/synchronized_videos"
+    # )
+
+    # synchronized_video_path = Path(
+    #     r'C:\Users\aaron\FreeMocap_Data\recording_sessions\freemocap_test_data\synchronized_videos'
+    # )
+    # tracker_name = "YOLOMediapipeComboTracker"
+    # num_processes = None
+
+    # process_folder_of_videos(
+    #     model_info=MediapipeModelInfo(),
+    #     tracking_params=MediapipeTrackingParams(),
+    #     synchronized_video_path=synchronized_video_path,
+    #     num_processes=num_processes,
+    # )
+
+    tracker_name = "OpenPoseTracker"
+    num_processes = 1
+
+    input_video_folder = Path(r'C:\Users\aaron\FreeMocap_Data\recording_sessions\freemocap_test_data')
+    input_video_filepath = input_video_folder/'synchronized_videos'
+
+    # output_video_folder = input_video_folder/'openpose_annotated_videos'
+    # output_video_folder.mkdir(parents=True, exist_ok=True)
+
+    output_json_path = input_video_folder/'output_data'/'raw_data'/'openpose_jsons'
+    output_json_path.mkdir(parents=True, exist_ok=True)
+
+    openpose_exe_path = r'C:\openpose'
 
     process_folder_of_videos(
-        tracker_name=tracker_name,
-        tracking_params=get_tracker_params(tracker_name=tracker_name),
-        synchronized_video_path=synchronized_video_path,
+        model_info=OpenPoseModelInfo(),
+        # tracker_name=tracker_name,
+        # tracking_params=get_tracker_params(tracker_name=tracker_name),
+        tracking_params=OpenPoseTrackingParams(
+           openpose_exe_path=str(openpose_exe_path),
+           output_json_path=str(output_json_path),
+        ),
+        synchronized_video_path=input_video_filepath,
         num_processes=num_processes,
     )

skellytracker/trackers/base_tracker/base_recorder.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 import logging
-from typing import Dict
+from pathlib import Path
+from typing import Dict, Union
 
 import numpy as np
 
@@ -44,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.
 
@@ -55,3 +56,19 @@ def save(self, file_path: str) -> None:
             self.process_tracked_objects()
         logger.info(f"Saving recorded objects to {file_path}")
         np.save(file_path, self.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"
+        )

skellytracker/trackers/base_tracker/base_tracker.py

@@ -159,3 +159,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] = None,
+        recorder: BaseRecorder = None,
+        **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."
+        )

skellytracker/trackers/openpose_tracker/openpose_model_info.py

@@ -0,0 +1,163 @@
+from skellytracker.trackers.base_tracker.base_tracking_params import BaseTrackingParams
+from skellytracker.trackers.base_tracker.model_info import ModelInfo
+
+
+
+class OpenPoseModelInfo(ModelInfo):
+    name = "openpose"
+    tracker_name = "OpenPoseTracker"
+
+
+    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"
+    ]
+
+
+
+    num_tracked_points = len(landmark_names)
+
+    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_exe_path: str
+    output_json_path: str
+    net_resolution: str = "-1x320"
+    number_people_max: int = 1
+    write_video: bool = True
+    openpose_output_resolution: str = "-1x-1"