Files

48 lines
1.4 KiB
Python

from __future__ import annotations
import cv2
import numpy as np
from app.exercises.dead_bug.types import DeadBugResult, Point
from app.vision.pose_types import _POSE_CONNECTIONS
def draw_landmarks(
image: np.ndarray,
landmarks: list[Point],
required_indices: tuple[int, ...],
connections: tuple[tuple[int, int], ...] | None = None,
visibility_threshold: float = 0.45,
line_color: tuple[int, int, int] = (65, 180, 255),
point_color: tuple[int, int, int] = (80, 255, 120),
line_thickness: int = 2,
point_radius: int = 4,
) -> None:
"""绘制人体骨架关键点与连接线(仅绘制可见度达标的点)"""
if connections is None:
connections = _POSE_CONNECTIONS
h, w = image.shape[:2]
for start, end in connections:
if start >= len(landmarks) or end >= len(landmarks):
continue
p1 = landmarks[start]
p2 = landmarks[end]
if p1.visibility < visibility_threshold or p2.visibility < visibility_threshold:
continue
cv2.line(
image,
(int(p1.x * w), int(p1.y * h)),
(int(p2.x * w), int(p2.y * h)),
line_color,
line_thickness,
)
for idx in required_indices:
if idx >= len(landmarks):
continue
p = landmarks[idx]
if p.visibility >= visibility_threshold:
cv2.circle(image, (int(p.x * w), int(p.y * h)), point_radius, point_color, -1)