Temp

src/python/pose_format/bin/pose_estimation.py

@@ -1,9 +1,7 @@
 #!/usr/bin/env python
 import argparse
 import os
-
 import cv2
-from pose_format.utils.holistic import load_holistic
 
 
 def load_video_frames(cap: cv2.VideoCapture):
@@ -15,7 +13,8 @@ def load_video_frames(cap: cv2.VideoCapture):
     cap.release()
 
 
-def pose_video(input_path: str, output_path: str, format: str, additional_config: dict = {'model_complexity': 1}, progress: bool = True):
+def pose_video(input_path: str, output_path: str, format: str, use_cpu: bool, additional_config: dict = {'model_complexity': 1}, progress: bool = True):
+
     # Load video frames
     print('Loading video ...')
     cap = cv2.VideoCapture(input_path)
@@ -27,18 +26,47 @@ def pose_video(input_path: str, output_path: str, format: str, additional_config
     # Perform pose estimation
     print('Estimating pose ...')
     if format == 'mediapipe':
+        from pose_format.utils.holistic import load_holistic
         pose = load_holistic(frames,
                              fps=fps,
                              width=width,
                              height=height,
                              progress=progress,
                              additional_holistic_config=additional_config)
+    elif format == 'openpose':
+        from pose_format.utils.openpose import estimate_and_load_openpose
+        pose = estimate_and_load_openpose(input_path, 
+                            output_path,   
+                            fps=fps, 
+                            width=width,
+                            height=height)
+    elif format == 'openpifpaf':
+        from pose_format.utils.openpifpaf import estimate_and_load_openpifpaf
+        pose = estimate_and_load_openpifpaf(frames,
+                            fps=fps, 
+                            use_cpu=use_cpu,
+                            width=width,
+                            height=height)
+    elif format == 'mmposewholebody':
+        from pose_format.utils.mmposewholebody import load_mmposewholebody
+        pose = load_mmposewholebody(input_path,
+                            output_path, 
+                            use_cpu=use_cpu,
+                            fps=fps,
+                            width=width,
+                            height=height)
     else:
         raise NotImplementedError('Pose format not supported')
 
+    # Build output filename
+    video_name = os.path.splitext(os.path.basename(input_path))[0]   # "SOUP" from ".../SOUP.mp4"
+    output_pose_file_path = os.path.join(output_path, f"{video_name}.pose")
+    os.makedirs(output_path, exist_ok=True)
+
     # Write
-    print('Saving to disk ...')
-    with open(output_path, "wb") as f:
+    print(f'Saving to disk at path {output_pose_file_path}')
+
+    with open(output_pose_file_path, "wb") as f:
         pose.write(f)
 
 
@@ -70,10 +98,14 @@ def main():
     parser.add_argument('-i', required=True, type=str, help='path to input video file')
     parser.add_argument('-o', required=True, type=str, help='path to output pose file')
     parser.add_argument('--format',
-                        choices=['mediapipe'],
+                        choices=['mediapipe', 'mmposewholebody', 'openpose', 'openpifpaf'],
                         default='mediapipe',
                         type=str,
                         help='type of pose estimation to use')
+    parser.add_argument('--usecpu', 
+                        default=False,
+                        type=bool, 
+                        help='option to use CPU for pose estimation (if supported)')
     parser.add_argument('--additional-config', type=str, help='additional configuration for the pose estimator')
 
     args = parser.parse_args()
@@ -82,7 +114,7 @@ def main():
         raise FileNotFoundError(f"Video file {args.i} not found")
 
     additional_config = parse_additional_config(args.additional_config)
-    pose_video(args.i, args.o, args.format, additional_config)
+    pose_video(args.i, args.o, args.format, args.usecpu, additional_config)
 
     # pip install . && video_to_pose -i como.mp4 -o como1.pose --format mediapipe
     # pip install . && video_to_pose -i como.mp4 -o como2.pose --format mediapipe --additional-config="model_complexity=2,smooth_landmarks=false,refine_face_landmarks=true"

src/python/pose_format/utils/cocowholebody133_header.py

@@ -0,0 +1,120 @@
+from ..pose_header import PoseHeaderComponent
+
+
+def cocowholebody_components():
+    """
+    Creates a list of COCO-Wholebody 133 components.
+
+    Returns
+    -------
+    list of PoseHeaderComponent
+        List of COCO-Wholebody 133 components.
+    """
+    BODY_POINTS = [
+        "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",
+        "left_big_toe","left_small_toe","left_heel",
+        "right_big_toe","right_small_toe","right_heel",
+    ]
+
+    FACE_POINTS = [f"face-{i}" for i in range(68)]
+    LEFT_HAND_POINTS = [f"left_hand_{i}" for i in range(21)]
+    RIGHT_HAND_POINTS = [f"right_hand_{i}" for i in range(21)]
+
+    def map_limbs(points, limbs):
+        index_map = {name: idx for idx, name in enumerate(points)}
+        return [
+            (index_map[a], index_map[b])
+            for (a, b) in limbs
+        ]
+
+    BODY_LIMBS_NAMES = [
+        ("left_ankle", "left_knee"),
+        ("left_knee", "left_hip"),
+        ("right_ankle", "right_knee"),
+        ("right_knee", "right_hip"),
+        ("left_hip", "right_hip"),
+        ("left_shoulder", "left_hip"),
+        ("right_shoulder", "right_hip"),
+        ("left_shoulder", "right_shoulder"),
+        ("left_shoulder", "left_elbow"),
+        ("right_shoulder", "right_elbow"),
+        ("left_elbow", "left_wrist"),
+        ("right_elbow", "right_wrist"),
+        ("left_eye", "right_eye"),
+        ("nose", "left_eye"),
+        ("nose", "right_eye"),
+        ("left_eye", "left_ear"),
+        ("right_eye", "right_ear"),
+        ("left_ear", "left_shoulder"),
+        ("right_ear", "right_shoulder"),
+        ("left_ankle", "left_big_toe"),
+        ("left_ankle", "left_small_toe"),
+        ("left_ankle", "left_heel"),
+        ("right_ankle", "right_big_toe"),
+        ("right_ankle", "right_small_toe"),
+        ("right_ankle", "right_heel"),
+    ]
+
+    LEFT_HAND_LIMBS_NAMES = [
+        ("left_hand_0", "left_hand_1"), ("left_hand_1", "left_hand_2"),
+        ("left_hand_2", "left_hand_3"), ("left_hand_3", "left_hand_4"),
+        ("left_hand_0", "left_hand_5"), ("left_hand_5", "left_hand_6"),
+        ("left_hand_6", "left_hand_7"), ("left_hand_7", "left_hand_8"),
+        ("left_hand_0", "left_hand_9"), ("left_hand_9", "left_hand_10"),
+        ("left_hand_10", "left_hand_11"), ("left_hand_11", "left_hand_12"),
+        ("left_hand_0", "left_hand_13"), ("left_hand_13", "left_hand_14"),
+        ("left_hand_14", "left_hand_15"), ("left_hand_15", "left_hand_16"),
+        ("left_hand_0", "left_hand_17"), ("left_hand_17", "left_hand_18"),
+        ("left_hand_18", "left_hand_19"), ("left_hand_19", "left_hand_20"),
+    ]
+
+    RIGHT_HAND_LIMBS_NAMES = [
+        ("right_hand_0", "right_hand_1"), ("right_hand_1", "right_hand_2"),
+        ("right_hand_2", "right_hand_3"), ("right_hand_3", "right_hand_4"),
+        ("right_hand_0", "right_hand_5"), ("right_hand_5", "right_hand_6"),
+        ("right_hand_6", "right_hand_7"), ("right_hand_7", "right_hand_8"),
+        ("right_hand_0", "right_hand_9"), ("right_hand_9", "right_hand_10"),
+        ("right_hand_10", "right_hand_11"), ("right_hand_11", "right_hand_12"),
+        ("right_hand_0", "right_hand_13"), ("right_hand_13", "right_hand_14"),
+        ("right_hand_14", "right_hand_15"), ("right_hand_15", "right_hand_16"),
+        ("right_hand_0", "right_hand_17"), ("right_hand_17", "right_hand_18"),
+        ("right_hand_18", "right_hand_19"), ("right_hand_19", "right_hand_20"),
+    ]
+
+    components = [
+        PoseHeaderComponent(
+            name="BODY",
+            points=BODY_POINTS,
+            limbs= map_limbs(BODY_POINTS, BODY_LIMBS_NAMES),  
+            colors=[(0,255,0)],
+            point_format="XYC"
+        ),
+
+        PoseHeaderComponent(
+            name="FACE",
+            points=FACE_POINTS,
+            limbs=[],   # WholeBody face mesh is huge, usually omitted
+            colors=[(255,255,255)],
+            point_format="XYC"
+        ),
+
+        PoseHeaderComponent(
+            name="LEFT_HAND",
+            points=LEFT_HAND_POINTS,
+            limbs= map_limbs(LEFT_HAND_POINTS, LEFT_HAND_LIMBS_NAMES), 
+            colors=[(0,255,0)],
+            point_format="XYC"
+        ),
+
+        PoseHeaderComponent(
+            name="RIGHT_HAND",
+            points=RIGHT_HAND_POINTS,
+            limbs= map_limbs(RIGHT_HAND_POINTS, RIGHT_HAND_LIMBS_NAMES),
+            colors=[(255,128,0)],
+            point_format="XYC"
+        ),
+    ]
+    return components

src/python/pose_format/utils/generic.py

@@ -13,7 +13,7 @@
 # from pose_format.utils.holistic import holistic_components
 # The import above creates an error: ImportError: Please install mediapipe with: pip install mediapipe
 
-KnownPoseFormat = Literal["holistic", "openpose", "openpose_135"]
+KnownPoseFormat = Literal["holistic", "openpose", "openpose_135", "coco_wholebody_133"]
 
 
 def get_component_names(
@@ -41,13 +41,22 @@ def detect_known_pose_format(pose_or_header: Union[Pose,PoseHeader]) -> KnownPos
 
     openpose_135_components = [c.name for c in OpenPose135_Components]
 
+    coco_wholebody_133_components = [
+        "BODY", 
+        "FACE", 
+        "LEFT_HAND", 
+        "RIGHT_HAND"
+    ]
+
     for component_name in component_names:
         if component_name in mediapipe_components:
             return "holistic"
         if component_name in openpose_components:
             return "openpose"
         if component_name in openpose_135_components:
             return "openpose_135"
+        if component_name in coco_wholebody_133_components:
+            return "coco_wholebody_133"
 
     raise ValueError(
         f"Could not detect pose format, unknown pose header schema with component names: {component_names}"
@@ -87,6 +96,9 @@ def pose_hide_legs(pose: Pose, remove: bool = False) -> Pose:
             # if any of the items in point_
         points_to_remove_dict = {"pose_keypoints_2d": point_names_to_remove}
 
+    elif known_pose_format == "coco_wholebody_133":
+        points_to_remove_dict = {"pose_keypoints_2d": list([str(i) for i in range(13, 23)])} # cocowholebody133 legs are points 13-22
+
     else:
         raise NotImplementedError(
             f"Unsupported pose header schema {known_pose_format} for {pose_hide_legs.__name__}: {pose.header}"
@@ -111,7 +123,9 @@ def pose_hide_legs(pose: Pose, remove: bool = False) -> Pose:
 
     return pose
 
-
+'''
+Returns a tuple of tuples containing the component and point names for the right and left shoulders.
+'''
 def pose_shoulders(pose_header: PoseHeader) -> Tuple[Tuple[str, str], Tuple[str, str]]:
     known_pose_format = detect_known_pose_format(pose_header)
 
@@ -123,6 +137,9 @@ def pose_shoulders(pose_header: PoseHeader) -> Tuple[Tuple[str, str], Tuple[str,
 
     if known_pose_format == "openpose":
         return ("pose_keypoints_2d", "RShoulder"), ("pose_keypoints_2d", "LShoulder")
+
+    if known_pose_format == "coco_wholebody_133":
+        return ("BODY", "right_shoulder"), ("BODY", "left_shoulder")
 
     raise NotImplementedError(
         f"Unsupported pose header schema {known_pose_format} for {pose_shoulders.__name__}: {pose_header}"
@@ -142,6 +159,12 @@ def hands_indexes(pose_header: PoseHeader)-> List[int]:
             pose_header.get_point_index("hand_left_keypoints_2d", "M_CMC"),
             pose_header.get_point_index("hand_right_keypoints_2d", "M_CMC"),
         ]
+    if known_pose_format == "coco_wholebody_133":
+        return [
+            #not sure this is correct, unclear what this is being used for
+            pose_header.get_point_index("LEFT_HAND", "left_hand_0"),
+            pose_header.get_point_index("RIGHT_HAND", "right_hand_0"),
+        ]
     raise NotImplementedError(
         f"Unsupported pose header schema {known_pose_format} for {hands_indexes.__name__}: {pose_header}"
     )
@@ -164,6 +187,12 @@ def hands_components(pose_header: PoseHeader)-> Tuple[Tuple[str, str], Tuple[str
     if known_pose_format == "openpose":
         return ("hand_left_keypoints_2d", "hand_right_keypoints_2d"), ("BASE", "P_CMC", "I_CMC"), ("BASE", "M_CMC")
 
+    if known_pose_format == "coco_wholebody_133":
+        # LEFT_HAND and RIGHT_HAND are component names
+        # Plane: points 0 (wrist), 5 (pinky base), 17 (index base)
+        # Line: points 0 (wrist) to 9 (middle finger)
+        return ("LEFT_HAND", "RIGHT_HAND"), ("0", "5", "17"), ("0", "9")
+
     raise NotImplementedError(
         f"Unsupported pose header schema '{known_pose_format}' for {hands_components.__name__}: {pose_header}"
     )
@@ -241,6 +270,8 @@ def get_hand_wrist_index(pose: Pose, hand: str)-> int:
         return pose.header.get_point_index(f"{hand.upper()}_HAND_LANDMARKS", "WRIST")
     if known_pose_format == "openpose":
         return pose.header.get_point_index(f"hand_{hand.lower()}_keypoints_2d", "BASE")
+    if known_pose_format == "coco_wholebody_133":
+        return pose.header.get_point_index(f"{hand.upper()}_HAND_POINTS", "0")  # wrist is point 0
     raise NotImplementedError(
         f"Unsupported pose header schema {known_pose_format} for {get_hand_wrist_index.__name__}: {pose.header}"
     )
@@ -252,6 +283,8 @@ def get_body_hand_wrist_index(pose: Pose, hand: str)-> int:
         return pose.header.get_point_index("POSE_LANDMARKS", f"{hand.upper()}_WRIST")
     if known_pose_format == "openpose":
         return pose.header.get_point_index("pose_keypoints_2d", f"{hand.upper()[0]}Wrist")
+    if known_pose_format == "coco_wholebody_133":
+        return pose.header.get_point_index("BODY_POINTS", f"{hand.lower()}_wrist")
     raise NotImplementedError(
         f"Unsupported pose header schema {known_pose_format} for {get_body_hand_wrist_index.__name__}: {pose.header}"
     )