Pose Estimation with MediaPipe

hip
  \
   knee        straight leg = 180 degrees
  /            bent to a right angle = 90 degrees
ankle

YOLO + tracker (lesson 26) -> crop each player -> pose on the crop

"""Joints, angles and simple movement rules from pose landmarks.
Works on any list of 33 landmarks with x, y, visibility (MediaPipe's NormalizedLandmark, or our pretend ones)."""
import math
 
import cv2
import numpy as np
 
# The numbers MediaPipe uses for its 33 landmarks (checked against mediapipe's own PoseLandmark list)
J = {"nose": 0, "l_shoulder": 11, "r_shoulder": 12, "l_elbow": 13, "r_elbow": 14, "l_wrist": 15, "r_wrist": 16,
     "l_hip": 23, "r_hip": 24, "l_knee": 25, "r_knee": 26, "l_ankle": 27, "r_ankle": 28}
 
BONES = [("l_shoulder", "r_shoulder"), ("l_shoulder", "l_elbow"), ("l_elbow", "l_wrist"),
         ("r_shoulder", "r_elbow"), ("r_elbow", "r_wrist"), ("l_shoulder", "l_hip"), ("r_shoulder", "r_hip"),
         ("l_hip", "r_hip"), ("l_hip", "l_knee"), ("l_knee", "l_ankle"), ("r_hip", "r_knee"), ("r_knee", "r_ankle")]
 
# Each angle is measured AT the middle joint, between the other two.
ANGLES = {"l_knee": ("l_hip", "l_knee", "l_ankle"), "r_knee": ("r_hip", "r_knee", "r_ankle"),
          "l_elbow": ("l_shoulder", "l_elbow", "l_wrist"), "r_elbow": ("r_shoulder", "r_elbow", "r_wrist")}
 
 
def point(lms, name, w, h, min_vis=0.5):
    """Landmark -> pixel position, or None if the model is not sure it can see that joint."""
    lm = lms[J[name]]
    vis = getattr(lm, "visibility", 1.0)
    if vis is None or vis < min_vis:
        return None
    return (lm.x * w, lm.y * h)                                  # normalized (0 to 1) -> pixels. Angles need PIXELS: see the lesson
 
 
def angle_at(a, b, c):
    """Angle in degrees at b, between the lines b->a and b->c. 180 = straight, 90 = right angle."""
    v1 = (a[0] - b[0], a[1] - b[1])
    v2 = (c[0] - b[0], c[1] - b[1])
    n1, n2 = math.hypot(*v1), math.hypot(*v2)
    if n1 == 0 or n2 == 0:
        return None
    cosang = (v1[0] * v2[0] + v1[1] * v2[1]) / (n1 * n2)
    return math.degrees(math.acos(max(-1.0, min(1.0, cosang))))
 
 
def joint_angle(lms, name, w, h, min_vis=0.5):
    a, b, c = (point(lms, n, w, h, min_vis) for n in ANGLES[name])
    if a is None or b is None or c is None:
        return None
    return angle_at(a, b, c)
 
 
class Smoother:
    """Exponential smoothing: new = alpha*measurement + (1-alpha)*old. Pose jitters frame to frame."""
 
    def __init__(self, alpha=0.4):
        self.alpha, self.value = alpha, None
 
    def update(self, x):
        if x is None:
            return self.value                                    # keep the last good value through a missed frame
        self.value = x if self.value is None else self.alpha * x + (1 - self.alpha) * self.value
        return self.value
 
 
def arm_raised(lms, w, h, side="r", margin=0.05, min_vis=0.5):
    """Wrist clearly above the shoulder (smaller y is higher in a picture)."""
    wr, sh = point(lms, f"{side}_wrist", w, h, min_vis), point(lms, f"{side}_shoulder", w, h, min_vis)
    if wr is None or sh is None:
        return None
    return wr[1] < sh[1] - margin * h
 
 
def leg_length(lms, w, h, side="r"):
    hip, knee, ank = (point(lms, f"{side}_{n}", w, h, 0.0) for n in ("hip", "knee", "ankle"))
    return math.dist(hip, knee) + math.dist(knee, ank)
 
 
def ankle_speed_in_legs_per_second(history, fps, w, h, side="r"):
    """history: list of landmark lists, newest last. Speed is divided by leg length, so it does not
    depend on how far from the camera the player is. 1.0 means one leg length per second."""
    if len(history) < 2:
        return 0.0
    a0 = point(history[-2], f"{side}_ankle", w, h, 0.0)
    a1 = point(history[-1], f"{side}_ankle", w, h, 0.0)
    return math.dist(a0, a1) * fps / max(leg_length(history[-1], w, h, side), 1e-6)
 
 
def draw_skeleton(frame, lms, min_vis=0.5):
    out = frame.copy()
    h, w = out.shape[:2]
    for a, b in BONES:
        pa, pb = point(lms, a, w, h, min_vis), point(lms, b, w, h, min_vis)
        if pa and pb:
            cv2.line(out, (int(pa[0]), int(pa[1])), (int(pb[0]), int(pb[1])), (0, 255, 255), 3)
    for name in J:
        p = point(lms, name, w, h, min_vis)
        if p:
            cv2.circle(out, (int(p[0]), int(p[1])), 5, (0, 0, 255), -1)
    return out

"""A pretend person made of sticks, so we can test the maths without a camera or a model.
The landmarks have the same fields as MediaPipe's (x, y, z, visibility), and the same 33 slots."""
import math
from types import SimpleNamespace
 
from skeleton import J
 
 
def blank():
    return [SimpleNamespace(x=0.0, y=0.0, z=0.0, visibility=0.0, presence=0.0) for _ in range(33)]
 
 
def put(lms, name, x, y, vis=1.0):
    lms[J[name]] = SimpleNamespace(x=x, y=y, z=0.0, visibility=vis, presence=1.0)
 
 
def make_person(r_knee_deg=180.0, r_wrist_up=False, cx=0.5, w=640, h=480, thigh=0.22, shin=0.22):
    """Standing person seen from the front, in a w x h picture. Lengths are fractions of the picture HEIGHT,
    so a 'real' leg of equal length in pixels has the same length in both directions."""
    lms = blank()
    sh_y, hip_y = 0.25, 0.50
    put(lms, "l_shoulder", cx - 0.06 * h / w, sh_y); put(lms, "r_shoulder", cx + 0.06 * h / w, sh_y)
    put(lms, "l_hip", cx - 0.04 * h / w, hip_y);     put(lms, "r_hip", cx + 0.04 * h / w, hip_y)
    put(lms, "l_knee", cx - 0.04 * h / w, hip_y + thigh); put(lms, "l_ankle", cx - 0.04 * h / w, hip_y + thigh + shin)
    # right leg: thigh straight down, shin rotated so the angle AT THE KNEE is r_knee_deg
    kx, ky = cx + 0.04 * h / w, hip_y + thigh
    put(lms, "r_knee", kx, ky)
    bend = math.radians(180.0 - r_knee_deg)                      # 0 = straight
    # the shin points down, rotated sideways by `bend` (x measured in pixels, converted back to fractions)
    put(lms, "r_ankle", kx + (math.sin(bend) * shin * h) / w, ky + math.cos(bend) * shin)
    put(lms, "l_elbow", cx - 0.12 * h / w, 0.38); put(lms, "l_wrist", cx - 0.13 * h / w, 0.5)
    put(lms, "r_elbow", cx + 0.12 * h / w, 0.38)
    put(lms, "r_wrist", cx + 0.13 * h / w, 0.12 if r_wrist_up else 0.5)
    return lms

import numpy as np
from fake_pose import blank, make_person, put
from skeleton import *
 
W, H = 640, 480
bad = 0
def check(name, ok, extra=""):
    global bad; bad += not ok; print("PASS" if ok else "FAIL", name, extra)
 
check("straight line is 180 degrees", abs(angle_at((0, 0), (1, 0), (2, 0)) - 180) < 1e-9)
check("right angle is 90 degrees", abs(angle_at((0, 1), (0, 0), (1, 0)) - 90) < 1e-9)
check("folded back is about 0 degrees", angle_at((1, 0), (0, 0), (1, 0.0001)) < 0.1)
check("zero-length arm returns None, no crash", angle_at((0, 0), (0, 0), (1, 1)) is None)
 
for deg in (180, 150, 120, 90, 60):
    got = joint_angle(make_person(r_knee_deg=deg), "r_knee", W, H)
    check(f"knee bent to {deg} degrees measured as {got:.1f}", abs(got - deg) < 0.5)
check("straight left leg is 180", abs(joint_angle(make_person(), "l_knee", W, H) - 180) < 0.5)
 
# the classic mistake: angles from the normalized numbers, not pixels
p = make_person(r_knee_deg=135)
hip, knee, ank = (p[J[n]] for n in ("r_hip", "r_knee", "r_ankle"))
wrong = angle_at((hip.x, hip.y), (knee.x, knee.y), (ank.x, ank.y))
check("using 0-1 numbers instead of pixels gives a WRONG angle", abs(wrong - 135) > 5, f"{wrong:.1f} instead of 135")
 
check("hidden joint gives None", joint_angle(make_person(), "r_knee", W, H, min_vis=1.1) is None)
lo = make_person(); lo[J["r_ankle"]].visibility = 0.2
check("low visibility ankle gives None", joint_angle(lo, "r_knee", W, H) is None)
 
check("arm down: not raised", arm_raised(make_person(r_wrist_up=False), W, H, "r") is False)
check("arm up: raised", arm_raised(make_person(r_wrist_up=True), W, H, "r") is True)
check("arm raised unknown if wrist hidden", arm_raised(blank(), W, H, "r") is None)
 
s = Smoother(0.5)
vals = [s.update(x) for x in [10, 10, 30, None, 30]]
check("smoother: first value kept, jump softened, None keeps last", vals[0] == 10 and 10 < vals[2] < 30 and vals[3] == vals[2], vals)
 
# kick: ankle speed in leg-lengths per second, independent of distance from the camera
fps = 30
still = [make_person(180) for _ in range(5)]
swing = [make_person(180), make_person(100)]
sp_still = ankle_speed_in_legs_per_second(still, fps, W, H)
sp_swing = ankle_speed_in_legs_per_second(swing, fps, W, H)
check("standing still: speed 0", sp_still == 0.0)
check("a quick knee snap is fast (more than 3 legs per second)", sp_swing > 3, f"{sp_swing:.1f}")
small = [make_person(180, thigh=0.11, shin=0.11), make_person(100, thigh=0.11, shin=0.11)]
sp_small = ankle_speed_in_legs_per_second(small, fps, W, H)
check("same kick, half the size (far away): about the same speed", abs(sp_small - sp_swing) / sp_swing < 0.1, f"{sp_small:.1f} vs {sp_swing:.1f}")
 
img = draw_skeleton(np.zeros((H, W, 3), np.uint8), make_person(r_wrist_up=True))
check("skeleton drawing changes the picture", img.sum() > 0)
print("failures:", bad)

"""Run MediaPipe's Pose Landmarker on a video and write the knee angles to a CSV.
Needs:  pip install mediapipe opencv-python
        and a model file, e.g. pose_landmarker_lite.task, downloaded from the Models section of Google's
        'Pose landmark detection guide for Python' (see the lesson for the link)."""
import csv
import os
 
import cv2
 
from skeleton import Smoother, draw_skeleton, joint_angle
 
 
def landmarks_from_result(result, person=0):
    """The first person's 33 landmarks, or None if nobody was found in this frame."""
    if result is None or not result.pose_landmarks or len(result.pose_landmarks) <= person:
        return None
    return result.pose_landmarks[person]
 
 
def make_landmarker(model_path, num_poses=1):
    if not os.path.isfile(model_path):
        raise FileNotFoundError(
            f"Model file not found: {model_path}\n"
            "Download pose_landmarker_lite.task (see the lesson) and put it next to this script.")
    import mediapipe as mp
    BaseOptions = mp.tasks.BaseOptions
    PoseLandmarker = mp.tasks.vision.PoseLandmarker
    PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
    RunningMode = mp.tasks.vision.RunningMode
    options = PoseLandmarkerOptions(base_options=BaseOptions(model_asset_path=model_path),
                                    running_mode=RunningMode.VIDEO, num_poses=num_poses)
    return PoseLandmarker.create_from_options(options)
 
 
def analyse_video(video_path, model_path, out_csv="knee_angles.csv", out_video=None):
    import mediapipe as mp
    landmarker = make_landmarker(model_path)
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    writer = cv2.VideoWriter(out_video, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) if out_video else None
    smooth_l, smooth_r = Smoother(0.4), Smoother(0.4)
    rows, i = [], 0
    with landmarker:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
            result = landmarker.detect_for_video(image, int(i * 1000 / fps))        # timestamps must increase every frame
            lms = landmarks_from_result(result)
            if lms is not None:
                left = smooth_l.update(joint_angle(lms, "l_knee", w, h))
                right = smooth_r.update(joint_angle(lms, "r_knee", w, h))
                if writer:
                    frame = draw_skeleton(frame, lms)
            else:
                left = right = None
            rows.append([i, round(i / fps, 2), None if left is None else round(left, 1), None if right is None else round(right, 1)])
            if writer:
                writer.write(frame)
            i += 1
    cap.release()
    if writer:
        writer.release()
    with open(out_csv, "w", newline="") as f:
        wr = csv.writer(f)
        wr.writerow(["frame", "t", "left_knee_deg", "right_knee_deg"])
        wr.writerows(rows)
    return len(rows)
 
 
if __name__ == "__main__":
    import sys
    print(analyse_video(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else "pose_landmarker_lite.task", out_video="pose_out.mp4"), "frames")

"""Checks the glue around MediaPipe using MediaPipe's own result classes. Needs mediapipe installed; no model needed."""
from mediapipe.tasks.python.components.containers.landmark import NormalizedLandmark
from mediapipe.tasks.python.vision import PoseLandmarkerResult
 
from pose_run import landmarks_from_result, make_landmarker
from skeleton import J, joint_angle
 
bad = 0
def check(name, ok, extra=""):
    global bad; bad += not ok; print("PASS" if ok else "FAIL", name, extra)
 
def result(n_people):
    people = [[NormalizedLandmark(x=0.5, y=0.5, z=0, visibility=1.0, presence=1.0) for _ in range(33)] for _ in range(n_people)]
    return PoseLandmarkerResult(pose_landmarks=people, pose_world_landmarks=[], segmentation_masks=None)
 
check("one person: 33 landmarks", len(landmarks_from_result(result(1))) == 33)
check("nobody found: None", landmarks_from_result(result(0)) is None)
check("no result at all: None", landmarks_from_result(None) is None)
check("asking for a second person who is not there: None", landmarks_from_result(result(1), person=1) is None)
try:
    make_landmarker("definitely_missing.task"); check("missing model gives a clear error", False)
except FileNotFoundError as e:
    check("missing model gives a clear error", "Download" in str(e))
 
# our joint numbers match MediaPipe's own
from mediapipe.tasks.python.vision import PoseLandmark
names = {"nose": "NOSE", "l_shoulder": "LEFT_SHOULDER", "r_shoulder": "RIGHT_SHOULDER", "l_elbow": "LEFT_ELBOW", "r_elbow": "RIGHT_ELBOW",
         "l_wrist": "LEFT_WRIST", "r_wrist": "RIGHT_WRIST", "l_hip": "LEFT_HIP", "r_hip": "RIGHT_HIP", "l_knee": "LEFT_KNEE",
         "r_knee": "RIGHT_KNEE", "l_ankle": "LEFT_ANKLE", "r_ankle": "RIGHT_ANKLE"}
check("our joint numbers equal MediaPipe's PoseLandmark values", all(getattr(PoseLandmark, v).value == J[k] for k, v in names.items()))
print("failures:", bad)