visualize.py

"""Overlay + display — OpenCV drawing only (no detection logic lives here).

Draws the skeleton, the red ``ID`` bounding box, and a dark info panel showing
live metrics, mirroring the example output in the framework diagram. Skeleton
connections are defined locally by MediaPipe landmark index so this module does
not need to import MediaPipe.
"""
from __future__ import annotations

from collections import deque
from typing import Deque, List, Optional, Tuple

import cv2

Landmark = Tuple[float, float, float]

# Landmarks below this visibility are faded out rather than drawn.
VIS_THRESHOLD = 0.5

# Body skeleton connections, by MediaPipe landmark index.
POSE_CONNECTIONS = [
    (11, 12),                      # shoulders
    (11, 13), (13, 15),            # left arm
    (12, 14), (14, 16),            # right arm
    (11, 23), (12, 24),            # torso sides
    (23, 24),                      # hips
    (23, 25), (25, 27), (27, 31),  # left leg + foot
    (24, 26), (26, 28), (28, 32),  # right leg + foot
]

_SKELETON_COLOR = (0, 255, 0)   # green
_JOINT_COLOR = (255, 255, 255)  # white
_BOX_COLOR = (0, 0, 255)        # red (BGR)
_PANEL_BG = (30, 30, 30)
_TEXT_COLOR = (255, 255, 255)
_TRAIL_COLOR = (255, 200, 0)    # cyan-ish (BGR)


class TrajectoryTrail:
    """A fading polyline of the body's recent centroid path.

    Holds the last ``maxlen`` centroid points in a deque; older points are drawn
    fainter and thinner so the trail trails off naturally.
    The centroid itself is computed in :mod:`analysis`.
    """

    def __init__(self, maxlen: int = 30):
        self._points: Deque[Tuple[int, int]] = deque(maxlen=maxlen)

    def update(self, point: Tuple[float, float]) -> None:
        self._points.append((int(point[0]), int(point[1])))

    def clear(self) -> None:
        self._points.clear()

    def draw(self, frame) -> None:
        pts = list(self._points)
        n = len(pts)
        for i in range(1, n):
            frac = i / n  # 0 (oldest) -> 1 (newest): fade in as the path nears now
            color = tuple(int(c * frac) for c in _TRAIL_COLOR)
            cv2.line(frame, pts[i - 1], pts[i], color, max(1, int(3 * frac)), cv2.LINE_AA)


def fit_to_display(frame, max_w: int = 1280, max_h: int = 720):
    """Return ``frame`` downscaled to fit within ``max_w`` x ``max_h``, preserving
    aspect ratio. Never upscales — small frames pass through unchanged.

    This is a display-only step: the pose pipeline still runs on the full-res
    frame, so landmark pixel coords, metrics, and overlays are unaffected.
    """
    h, w = frame.shape[:2]
    scale = min(max_w / w, max_h / h, 1.0)
    if scale == 1.0:
        return frame
    new_size = (int(w * scale), int(h * scale))
    return cv2.resize(frame, new_size, interpolation=cv2.INTER_AREA)


def draw_skeleton(frame, landmarks: List[Landmark]) -> None:
    for a, b in POSE_CONNECTIONS:
        pa, pb = landmarks[a], landmarks[b]
        if pa[2] < VIS_THRESHOLD or pb[2] < VIS_THRESHOLD:
            continue
        cv2.line(
            frame,
            (int(pa[0]), int(pa[1])),
            (int(pb[0]), int(pb[1])),
            _SKELETON_COLOR,
            4,
        )
    for x, y, v in landmarks:
        if v < VIS_THRESHOLD:
            continue
        cv2.circle(frame, (int(x), int(y)), 6, _JOINT_COLOR, -1)


def draw_id_box(frame, landmarks: List[Landmark], track_id: Optional[int]) -> None:
    if track_id is None:
        return
    xs = [x for x, _, v in landmarks if v >= VIS_THRESHOLD]
    ys = [y for _, y, v in landmarks if v >= VIS_THRESHOLD]
    if not xs or not ys:
        return
    pad = 15
    x1, y1 = max(0, int(min(xs)) - pad), max(0, int(min(ys)) - pad)
    x2, y2 = int(max(xs)) + pad, int(max(ys)) + pad
    cv2.rectangle(frame, (x1, y1), (x2, y2), _BOX_COLOR, 4)
    cv2.putText(
        frame, f"ID: {track_id}", (x1, max(12, y1 - 8)),
        cv2.FONT_HERSHEY_SIMPLEX, 0.7, _BOX_COLOR, 2, cv2.LINE_AA,
    )


def draw_info_panel(frame, metrics: dict, fps: Optional[float] = None) -> None:
    lines = []
    kl = metrics.get("knee_angle_left")
    lines.append(f"L Knee: {kl:.0f} deg" if kl is not None else "L Knee: --")
    kr = metrics.get("knee_angle_right")
    lines.append(f"R Knee: {kr:.0f} deg" if kr is not None else "R Knee: --")
    sp = metrics.get("speed_px_s")
    lines.append(f"Speed: {sp:.0f} px/s" if sp is not None else "Speed: --")
    activity = metrics.get("activity")
    lines.append(f"Activity: {activity}" if activity else "Activity: --")
    # Reps follow the recognised activity (squat/jumping jacks/push-ups); a
    # locomotion state reports 0, which we render as "--".
    reps = metrics.get("activity_reps")
    lines.append(f"Reps: {reps}" if reps else "Reps: --")
    if fps is not None:
        lines.append(f"FPS: {fps:.0f}")

    pad, line_h, width = 10, 26, 230
    height = pad * 2 + line_h * len(lines)
    overlay = frame.copy()
    cv2.rectangle(overlay, (10, 10), (10 + width, 10 + height), _PANEL_BG, -1)
    cv2.addWeighted(overlay, 0.6, frame, 0.4, 0, frame)

    y = 10 + pad + 18
    for line in lines:
        cv2.putText(
            frame, line, (20, y),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, _TEXT_COLOR, 1, cv2.LINE_AA,
        )
        y += line_h