wehub-resource-sync
264 lines
9.5 KiB
Python
Executable File
264 lines
9.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Use the MediaPipe Pose solution to detect and track a human pose in video."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
from contextlib import closing
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Final
|
|
|
|
import cv2
|
|
import mediapipe as mp
|
|
import mediapipe.python.solutions.pose as mp_pose
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import requests
|
|
|
|
import rerun as rr # pip install rerun-sdk
|
|
import rerun.blueprint as rrb
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator
|
|
|
|
DESCRIPTION = """
|
|
# Human pose tracking
|
|
This example uses Rerun to visualize the output of [MediaPipe](https://developers.google.com/mediapipe)-based tracking
|
|
of a human pose in 2D and 3D.
|
|
|
|
The full source code for this example is available
|
|
[on GitHub](https://github.com/rerun-io/rerun/blob/latest/examples/python/human_pose_tracking).
|
|
""".strip()
|
|
|
|
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
|
|
DATASET_DIR: Final = EXAMPLE_DIR / "dataset" / "pose_movement"
|
|
MODEL_DIR: Final = EXAMPLE_DIR / "model" / "pose_movement"
|
|
DATASET_URL_BASE: Final = "https://storage.googleapis.com/rerun-example-datasets/pose_movement"
|
|
MODEL_URL_TEMPLATE: Final = "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_{model_name}/float16/latest/pose_landmarker_{model_name}.task"
|
|
|
|
|
|
def track_pose(video_path: str, model_path: str, *, max_frame_count: int | None) -> None:
|
|
options = mp.tasks.vision.PoseLandmarkerOptions(
|
|
base_options=mp.tasks.BaseOptions(
|
|
model_asset_path=model_path,
|
|
),
|
|
running_mode=mp.tasks.vision.RunningMode.VIDEO,
|
|
output_segmentation_masks=True,
|
|
)
|
|
|
|
rr.log("description", rr.TextDocument(DESCRIPTION, media_type=rr.MediaType.MARKDOWN), static=True)
|
|
|
|
rr.log(
|
|
"/",
|
|
rr.AnnotationContext(
|
|
rr.ClassDescription(
|
|
info=rr.AnnotationInfo(id=1, label="Person"),
|
|
keypoint_annotations=[rr.AnnotationInfo(id=lm.value, label=lm.name) for lm in mp_pose.PoseLandmark],
|
|
keypoint_connections=mp_pose.POSE_CONNECTIONS,
|
|
),
|
|
),
|
|
static=True,
|
|
)
|
|
# Use a separate annotation context for the segmentation mask.
|
|
rr.log(
|
|
"video/mask",
|
|
rr.AnnotationContext([
|
|
rr.AnnotationInfo(id=0, label="Background"),
|
|
rr.AnnotationInfo(id=1, label="Person", color=(0, 0, 0)),
|
|
]),
|
|
static=True,
|
|
)
|
|
rr.log("person", rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, static=True)
|
|
|
|
pose_landmarker = mp.tasks.vision.PoseLandmarker.create_from_options(options)
|
|
|
|
with closing(VideoSource(video_path)) as video_source:
|
|
for idx, bgr_frame in enumerate(video_source.stream_bgr()):
|
|
if max_frame_count is not None and idx >= max_frame_count:
|
|
break
|
|
|
|
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=bgr_frame.data)
|
|
rr.set_time("time", duration=bgr_frame.time)
|
|
rr.set_time("frame_idx", sequence=bgr_frame.idx)
|
|
|
|
results = pose_landmarker.detect_for_video(mp_image, int(bgr_frame.time * 1000))
|
|
h, w, _ = bgr_frame.data.shape
|
|
landmark_positions_2d = read_landmark_positions_2d(results, w, h)
|
|
|
|
rr.log("video/bgr", rr.Image(bgr_frame.data, color_model="BGR").compress(jpeg_quality=75))
|
|
if landmark_positions_2d is not None:
|
|
rr.log(
|
|
"video/pose/points",
|
|
rr.Points2D(landmark_positions_2d, class_ids=1, keypoint_ids=mp_pose.PoseLandmark),
|
|
)
|
|
|
|
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=1, keypoint_ids=mp_pose.PoseLandmark),
|
|
)
|
|
|
|
if results.segmentation_masks is not None:
|
|
segmentation_mask = results.segmentation_masks[0].numpy_view()
|
|
binary_segmentation_mask = segmentation_mask > 0.5
|
|
rr.log("video/mask", rr.SegmentationImage(binary_segmentation_mask.astype(np.uint8)))
|
|
|
|
|
|
def read_landmark_positions_2d(
|
|
results: Any,
|
|
image_width: int,
|
|
image_height: int,
|
|
) -> npt.NDArray[np.float32] | None:
|
|
if results.pose_landmarks is None or len(results.pose_landmarks) == 0:
|
|
return None
|
|
else:
|
|
pose_landmarks = results.pose_landmarks[0]
|
|
normalized_landmarks = [pose_landmarks[lm] for lm in mp_pose.PoseLandmark]
|
|
return np.array([(image_width * lm.x, image_height * lm.y) for lm in normalized_landmarks])
|
|
|
|
|
|
def read_landmark_positions_3d(
|
|
results: Any,
|
|
) -> npt.NDArray[np.float32] | None:
|
|
if results.pose_landmarks is None or len(results.pose_landmarks) == 0:
|
|
return None
|
|
else:
|
|
pose_landmarks = results.pose_landmarks[0]
|
|
landmarks = [pose_landmarks[lm] for lm in mp_pose.PoseLandmark]
|
|
return np.array([(lm.x, lm.y, lm.z) for lm in landmarks])
|
|
|
|
|
|
@dataclass
|
|
class VideoFrame:
|
|
data: cv2.typing.MatLike
|
|
time: float
|
|
idx: int
|
|
|
|
|
|
class VideoSource:
|
|
def __init__(self, path: str) -> None:
|
|
self.capture = cv2.VideoCapture(path)
|
|
|
|
if not self.capture.isOpened():
|
|
logging.error("Couldn't open video at %s", path)
|
|
|
|
def close(self) -> None:
|
|
self.capture.release()
|
|
|
|
def stream_bgr(self) -> Iterator[VideoFrame]:
|
|
while self.capture.isOpened():
|
|
idx = int(self.capture.get(cv2.CAP_PROP_POS_FRAMES))
|
|
is_open, bgr = self.capture.read()
|
|
time_ms = self.capture.get(cv2.CAP_PROP_POS_MSEC)
|
|
|
|
if not is_open:
|
|
break
|
|
|
|
yield VideoFrame(data=bgr, time=time_ms * 1e-3, idx=idx)
|
|
|
|
|
|
def get_downloaded_video_path(dataset_dir: Path, video_name: str) -> str:
|
|
video_file_name = f"{video_name}.mp4"
|
|
destination_path = dataset_dir / video_file_name
|
|
if destination_path.exists():
|
|
logging.info("%s already exists. No need to download", destination_path)
|
|
return str(destination_path)
|
|
|
|
source_path = f"{DATASET_URL_BASE}/{video_file_name}"
|
|
|
|
logging.info("Downloading video from %s to %s", source_path, destination_path)
|
|
os.makedirs(dataset_dir.absolute(), exist_ok=True)
|
|
download(source_path, destination_path)
|
|
return str(destination_path)
|
|
|
|
|
|
def get_downloaded_model_path(model_dir: Path, model_name: str) -> str:
|
|
model_file_name = f"{model_name}.task"
|
|
destination_path = model_dir / model_file_name
|
|
if destination_path.exists():
|
|
logging.info("%s already exists. No need to download", destination_path)
|
|
return str(destination_path)
|
|
|
|
model_url = MODEL_URL_TEMPLATE.format(model_name=model_name)
|
|
logging.info("Downloading model from %s to %s", model_url, destination_path)
|
|
download(model_url, destination_path)
|
|
return str(destination_path)
|
|
|
|
|
|
def download(url: str, destination_path: Path) -> None:
|
|
os.makedirs(destination_path.parent, exist_ok=True)
|
|
with requests.get(url, stream=True) as req:
|
|
req.raise_for_status()
|
|
with open(destination_path, "wb") as f:
|
|
f.writelines(req.iter_content(chunk_size=8192))
|
|
|
|
|
|
def main() -> None:
|
|
# Ensure the logging gets written to stderr:
|
|
logging.getLogger().addHandler(logging.StreamHandler())
|
|
logging.getLogger().setLevel("INFO")
|
|
|
|
parser = argparse.ArgumentParser(description="Uses the MediaPipe Pose solution to track a human pose in video.")
|
|
parser.add_argument(
|
|
"--video",
|
|
type=str,
|
|
default="backflip",
|
|
choices=["backflip", "soccer"],
|
|
help="The example video to run on.",
|
|
)
|
|
parser.add_argument("--dataset-dir", type=Path, default=DATASET_DIR, help="Directory to save example videos to.")
|
|
parser.add_argument("--video-path", type=str, default="", help="Full path to video to run on. Overrides `--video`.")
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="heavy",
|
|
choices=["lite", "full", "heavy"],
|
|
help="The mediapipe model to use (see https://developers.google.com/mediapipe/solutions/vision/pose_landmarker).",
|
|
)
|
|
parser.add_argument("--model-dir", type=Path, default=MODEL_DIR, help="Directory to save downloaded model to.")
|
|
parser.add_argument("--model-path", type=str, default="", help="Full path of mediapipe model. Overrides `--model`.")
|
|
parser.add_argument(
|
|
"--max-frame",
|
|
type=int,
|
|
help="Stop after processing this many frames. If not specified, will run until interrupted.",
|
|
)
|
|
rr.script_add_args(parser)
|
|
|
|
args = parser.parse_args()
|
|
rr.script_setup(
|
|
args,
|
|
"rerun_example_human_pose_tracking",
|
|
default_blueprint=rrb.Horizontal(
|
|
rrb.Vertical(
|
|
rrb.Spatial2DView(origin="video", name="Result"),
|
|
rrb.Spatial3DView(origin="person", name="3D pose"),
|
|
),
|
|
rrb.Vertical(
|
|
rrb.Spatial2DView(origin="video/bgr", name="Raw video"),
|
|
rrb.TextDocumentView(origin="description", name="Description"),
|
|
row_shares=[2, 3],
|
|
),
|
|
column_shares=[3, 2],
|
|
),
|
|
)
|
|
|
|
video_path = args.video_path # type: str
|
|
if not video_path:
|
|
video_path = get_downloaded_video_path(args.dataset_dir, args.video)
|
|
|
|
model_path = args.model_path # type: str
|
|
if not args.model_path:
|
|
model_path = get_downloaded_model_path(args.model_dir, args.model)
|
|
|
|
track_pose(video_path, model_path, max_frame_count=args.max_frame)
|
|
|
|
rr.script_teardown(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|