# 安装所需的Python包
pip install mediapipe
pip install numpy
pip install opencv-python<4.6
pip install requests>=2.31,<3
pip install rerun-sdk

# 也可以直接使用配置文件requirements.txt
pip install -r examples/python/human_pose_tracking/requirements.txt

import mediapipe as mp
import numpy as np
from typing import Any
import numpy.typing as npt
import cv2


"""

Args:
 results (Any): Mediapipe Pose results.
 image_width (int): Width of the input image.
 image_height (int): Height of the input image.

 Returns:
 np.array | None: Array of 2D landmark positions or None if no landmarks are detected.
"""
def read_landmark_positions_2d(
 results: Any,
 image_width: int,
 image_height: int,
) -> npt.NDArray[np.float32] | None:
 if results.pose_landmarks is None:
 return None
 else:
 # 提取标准化的地标位置并将其缩放为图像尺寸
 normalized_landmarks = [results.pose_landmarks.landmark[lm] for lm in mp.solutions.pose.PoseLandmark]
 return np.array([(image_width * lm.x, image_height * lm.y) for lm in normalized_landmarks])


"""

Args:
 results (Any): Mediapipe Pose results.

 Returns:
 np.array | None: Array of 3D landmark positions or None if no landmarks are detected.
"""
def read_landmark_positions_3d(
 results: Any,
) -> npt.NDArray[np.float32] | None:
 if results.pose_landmarks is None:
 return None
 else:
 # 提取三维地标位置
 landmarks = [results.pose_world_landmarks.landmark[lm] for lm in mp.solutions.pose.PoseLandmark]
 return np.array([(lm.x, lm.y, lm.z) for lm in landmarks])


"""

Args:
 image_path (str): Path to the input image.
"""
def track_pose(image_path: str) -> None:
 # 读取图像，将颜色转换为RGB格式
 image = cv2.imread(image_path)
 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

 # 创建Pose模型实例
 pose_detector = mp.solutions.pose.Pose(static_image_mode=True)

 # 处理图像以获得姿势标志
 results = pose_detector.process(image)
 h, w, _ = image.shape

 # 读取二维和三维地标位置
 landmark_positions_2d = read_landmark_positions_2d(results, w, h)
 landmark_positions_3d = read_landmark_positions_3d(results)

rr.log(
 "/",
 rr.AnnotationContext(
 rr.ClassDescription(
 info=rr.AnnotationInfo(id=0, label="Person"),
 keypoint_annotatinotallow=[rr.AnnotationInfo(id=lm.value, label=lm.name) for lm in mp_pose.PoseLandmark],
 keypoint_cnotallow=mp_pose.POSE_CONNECTIONS,
 )
 ),
 timeless=True,
 )

def track_pose_2d(video_path: str) -> None:
 mp_pose = mp.solutions.pose 

 with closing(VideoSource(video_path)) as video_source, mp_pose.Pose() as pose:
 for idx, bgr_frame in enumerate(video_source.stream_bgr()):
 if max_frame_count is not None and idx >= max_frame_count:
 break

 rgb = cv2.cvtColor(bgr_frame.data, cv2.COLOR_BGR2RGB)

 # 将帧与数据关联
 rr.set_time_seconds("time", bgr_frame.time)
 rr.set_time_sequence("frame_idx", bgr_frame.idx)

 # 呈现视频
 rr.log("video/rgb", rr.Image(rgb).compress(jpeg_quality=75))

 # 获取预测结果
 results = pose.process(rgb)
 h, w, _ = rgb.shape

 # 把2D点记录到'video'实体中
 landmark_positions_2d = read_landmark_positions_2d(results, w, h)
 if landmark_positions_2d is not None:
 rr.log(
 "video/pose/points",
 rr.Points2D(landmark_positions_2d, class_ids=0, keypoint_ids=mp_pose.PoseLandmark),
 )

def track_pose_3d(video_path: str, *, segment: bool, max_frame_count: int | None) -> None:
 mp_pose = mp.solutions.pose 

 rr.log("person", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, timeless=True)

 with closing(VideoSource(video_path)) as video_source, mp_pose.Pose() as pose:
 for idx, bgr_frame in enumerate(video_source.stream_bgr()):
 if max_frame_count is not None and idx >= max_frame_count:
 break

 rgb = cv2.cvtColor(bgr_frame.data, cv2.COLOR_BGR2RGB)

 # 把帧与数据关联起来
 rr.set_time_seconds("time", bgr_frame.time)
 rr.set_time_sequence("frame_idx", bgr_frame.idx)

 # 呈现视频
 rr.log("video/rgb", rr.Image(rgb).compress(jpeg_quality=75))

 # 取得预测结果
 results = pose.process(rgb)
 h, w, _ = rgb.shape

 # 对于3D呈现的新的实例"Person"
 landmark_positions_3d = read_landmark_positions_3d(results)
 if landmark_positions_3d is not None:
 rr.log(
 "person/pose/points",
 rr.Points3D(landmark_positions_3d, class_ids=0, keypoint_ids=mp_pose.PoseLandmark),
 )

rr.log(
 "video", 
 rr.Image(img).compress(jpeg_quality=75)
)