Spaces:

Stouching-Robotics
/

DataFlowing

Running

DataFlowing / processor /app /skeleton_renderer.py

stouchinglhh

collector: 同步采集工程 v1.2.1（全部 LFS 模式文件指针化，与 Space 现有文件合并）

d646dd7 2 months ago

10.9 kB

"""Hand skeleton video pre-rendering using Google MediaPipe drawing_utils.
Reads a video + parquet hand keypoints, renders skeletons onto every frame
using official MediaPipe styling, and writes a new MP4 alongside the original.
Usage::
    from app.skeleton_renderer import render_skeleton_video
    output = await render_skeleton_video(video_path, session_dir)
    # → video_path.parent / "video_skeleton.mp4"  (or None if no hand data)
"""

from __future__ import annotations

import logging
import subprocess
from pathlib import Path

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# ═══════════════════════════════════════════════════════════════
# Helpers
# ═══════════════════════════════════════════════════════════════

def _has_hand_keypoints(parquet_path: Path) -> bool:
    """Check whether a merged data parquet has hand keypoint columns with data."""
    try:
        df = pd.read_parquet(parquet_path, columns=["hand_0_keypoints", "hand_1_keypoints"])
    except (ValueError, Exception):
        return False

    for col in ("hand_0_keypoints", "hand_1_keypoints"):
        if col in df.columns:
            has_data = df[col].notna().any()
            if has_data:
                return True
    return False


def _find_data_parquet(session_dir: Path) -> Path | None:
    """Find the merged data parquet in a session directory."""
    for d in session_dir.rglob("data"):
        if not d.is_dir():
            continue
        parq_files = sorted(
            p for p in d.rglob("chunk_*.parquet")
            if "meta" not in str(p)
            and not p.name.startswith("auto_labels_")
            and not p.name.startswith("hand_kpts_")
        )
        if parq_files:
            return parq_files[0]
    return None


def _build_keypoint_lookup(parquet_path: Path) -> dict[int, dict]:
    """Build a frame_index → hand keypoints lookup dict.
    Returns dict like::
        {0: {"h0": [[x,y]*21], "h1": [[x,y]*21]}, 30: {"h0": [...], "h1": None}, ...}
    """
    df = pd.read_parquet(parquet_path)
    lookup: dict[int, dict] = {}

    has_h0 = "hand_0_keypoints" in df.columns
    has_h1 = "hand_1_keypoints" in df.columns

    if not has_h0 and not has_h1:
        return lookup

    for _, row in df.iterrows():
        fi = int(row["frame_index"])
        entry: dict = {}
        for hk, col in (("h0", "hand_0_keypoints"), ("h1", "hand_1_keypoints")):
            if col not in df.columns:
                continue
            val = row[col]
            if val is None or (isinstance(val, float) and pd.isna(val)):
                entry[hk] = None
                continue
            kp_list = val.tolist() if hasattr(val, "tolist") else val
            if isinstance(kp_list, list) and len(kp_list) >= 21:
                entry[hk] = [[float(p[0]), float(p[1])] for p in kp_list[:21]]
            else:
                entry[hk] = None
        lookup[fi] = entry

    return lookup


# ═══════════════════════════════════════════════════════════════
# Drawing
# ═══════════════════════════════════════════════════════════════

def _get_drawing_func():
    """Return the best available drawing function.
    Tries Google MediaPipe official drawing_utils first,
    falls back to OpenCV-based rendering.
    """
    try:
        from mediapipe.tasks.python.vision import drawing_utils, drawing_styles
        from mediapipe.tasks.python.vision.hand_landmarker import HandLandmarksConnections
        from mediapipe.tasks.python.components.containers import landmark as lm

        def _draw_mediapipe(frame_bgr, kp_lookup, frame_index):
            """Draw hand skeletons on a BGR frame using Google MediaPipe official API."""
            entry = kp_lookup.get(frame_index)
            if not entry:
                return

            W, H = 640, 480  # source resolution

            for hk in ("h0", "h1"):
                kp = entry.get(hk)
                if not kp or len(kp) < 21:
                    continue

                landmarks = []
                for pt in kp[:21]:
                    l = lm.NormalizedLandmark()
                    l.x = pt[0] / W
                    l.y = pt[1] / H
                    l.z = 0.0
                    landmarks.append(l)

                drawing_utils.draw_landmarks(
                    frame_bgr,
                    landmarks,
                    HandLandmarksConnections.HAND_CONNECTIONS,
                    drawing_styles.get_default_hand_landmarks_style(),
                    drawing_styles.get_default_hand_connections_style(),
                )

        return _draw_mediapipe, "mediapipe"

    except ImportError:
        pass

    # ── OpenCV fallback ──
    def _draw_opencv(frame_bgr, kp_lookup, frame_index):
        """Draw hand skeletons using OpenCV (near-MediaPipe replica)."""
        entry = kp_lookup.get(frame_index)
        if not entry:
            return

        import cv2

        connections = [
            (0,1),(1,2),(2,3),(3,4),(0,5),(5,6),(6,7),(7,8),
            (0,9),(9,10),(10,11),(11,12),(0,13),(13,14),(14,15),(15,16),
            (0,17),(17,18),(18,19),(19,20),(5,9),(9,13),(13,17),
        ]
        # BGR colours per finger (matches MediaPipe style)
        finger_colors = [
            (255, 48, 48),   # thumb  → red
            (128, 64, 128),  # index  → purple
            (48, 255, 255),  # middle → yellow
            (48, 255, 48),   # ring   → green
            (48, 208, 255),  # pinky  → orange
        ]

        for hk in ("h0", "h1"):
            kp = entry.get(hk)
            if not kp or len(kp) < 21:
                continue
            pts = [(int(p[0]), int(p[1])) for p in kp[:21]]

            # Palm connections (gray)
            for i, j in [(0, 1), (0, 17), (1, 5), (5, 9), (9, 13), (13, 17)]:
                cv2.line(frame_bgr, pts[i], pts[j], (128, 128, 128), 3, cv2.LINE_AA)

            # Finger connections (colored)
            for f in range(5):
                c = finger_colors[f]
                base = f * 4 + 1
                for i in range(3):
                    cv2.line(frame_bgr, pts[base + i], pts[base + i + 1],
                             c, 2, cv2.LINE_AA)

            # Landmark dots: white border + colored fill
            radii = [7, 4, 4, 3, 2, 4, 4, 3, 2, 4, 4, 3, 2,
                     4, 4, 3, 2, 4, 4, 3, 2]
            for i, (x, y) in enumerate(pts):
                if i == 0:
                    fc = (255, 255, 255)  # wrist white
                else:
                    fc = finger_colors[(i - 1) // 4]
                r = radii[i]
                cv2.circle(frame_bgr, (x, y), r + 1, (255, 255, 255), -1, cv2.LINE_AA)
                cv2.circle(frame_bgr, (x, y), max(1, r - 1), fc, -1, cv2.LINE_AA)

    return _draw_opencv, "opencv"


# ═══════════════════════════════════════════════════════════════
# Main entry point
# ═══════════════════════════════════════════════════════════════

async def render_skeleton_video(
    video_path: Path,
    session_dir: Path,
) -> Path | None:
    """Render hand skeletons onto a video, returning the output path.
    Args:
        video_path: Path to the source MP4 file.
        session_dir: Session root directory (contains data/ with parquet).
    Returns:
        Path to the skeleton video, or ``None`` if no hand data exists.
    """
    import cv2

    # 1. Find and validate hand data
    parquet_path = _find_data_parquet(session_dir)
    if parquet_path is None:
        logger.info("No data parquet found in %s — skipping skeleton render", session_dir)
        return None

    if not _has_hand_keypoints(parquet_path):
        logger.info("No hand keypoint columns in %s — skipping skeleton render", parquet_path.name)
        return None

    # 2. Build keypoint lookup (frame_index → keypoints)
    kp_lookup = _build_keypoint_lookup(parquet_path)
    keypoint_frames = sum(1 for v in kp_lookup.values() if v.get("h0") or v.get("h1"))
    if keypoint_frames == 0:
        logger.info("No frames with hand keypoints — skipping skeleton render")
        return None

    logger.info(
        "Rendering skeleton video: %d / %d total rows have hand keypoints",
        keypoint_frames, len(kp_lookup),
    )

    # 3. Get drawing function
    draw_func, backend = _get_drawing_func()
    logger.info("Using drawing backend: %s", backend)

    # 4. Open video
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        logger.error("Cannot open video: %s", video_path)
        return None

    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps <= 0:
        fps = 30.0
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    logger.info("Video: %dx%d @ %.1f fps, %d frames", width, height, fps, total_frames)

    # 5. Output path: {stem}_skeleton{ext}
    output_path = video_path.parent / f"{video_path.stem}_skeleton{video_path.suffix}"

    # 6. FFmpeg H.264 encoder (browser-compatible)
    ffmpeg_cmd = [
        "ffmpeg", "-y",
        "-f", "rawvideo", "-vcodec", "rawvideo",
        "-s", f"{width}x{height}", "-pix_fmt", "bgr24", "-r", str(fps),
        "-i", "pipe:0",
        "-c:v", "libx264", "-preset", "fast", "-crf", "23",
        "-pix_fmt", "yuv420p",
        str(output_path),
    ]
    proc = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE, stderr=subprocess.DEVNULL)
    frame_size = width * height * 3  # BGR24 = 3 bytes per pixel

    # 7. Process frames
    frame_idx = 0
    rendered_frames = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        draw_func(frame, kp_lookup, frame_idx)
        proc.stdin.write(frame.tobytes())

        if kp_lookup.get(frame_idx):
            rendered_frames += 1

        frame_idx += 1
        if frame_idx % 50 == 0:
            logger.debug("Skeleton render: %d / %d frames", frame_idx, total_frames)

    cap.release()
    proc.stdin.close()
    proc.wait()

    logger.info(
        "Skeleton video written: %s (%d frames rendered, %d total)",
        output_path.name, rendered_frames, frame_idx,
    )
    return output_path

	"""Hand skeleton video pre-rendering using Google MediaPipe drawing_utils.

	Reads a video + parquet hand keypoints, renders skeletons onto every frame
	using official MediaPipe styling, and writes a new MP4 alongside the original.

	Usage::

	from app.skeleton_renderer import render_skeleton_video
	output = await render_skeleton_video(video_path, session_dir)
	# → video_path.parent / "video_skeleton.mp4" (or None if no hand data)
	"""

	from __future__ import annotations

	import logging
	import subprocess
	from pathlib import Path

	import numpy as np
	import pandas as pd

	logger = logging.getLogger(__name__)

	# ═══════════════════════════════════════════════════════════════
	# Helpers
	# ═══════════════════════════════════════════════════════════════

	def _has_hand_keypoints(parquet_path: Path) -> bool:
	"""Check whether a merged data parquet has hand keypoint columns with data."""
	try:
	df = pd.read_parquet(parquet_path, columns=["hand_0_keypoints", "hand_1_keypoints"])
	except (ValueError, Exception):
	return False

	for col in ("hand_0_keypoints", "hand_1_keypoints"):
	if col in df.columns:
	has_data = df[col].notna().any()
	if has_data:
	return True
	return False


	def _find_data_parquet(session_dir: Path) -> Path \| None:
	"""Find the merged data parquet in a session directory."""
	for d in session_dir.rglob("data"):
	if not d.is_dir():
	continue
	parq_files = sorted(
	p for p in d.rglob("chunk_*.parquet")
	if "meta" not in str(p)
	and not p.name.startswith("auto_labels_")
	and not p.name.startswith("hand_kpts_")
	)
	if parq_files:
	return parq_files[0]
	return None


	def _build_keypoint_lookup(parquet_path: Path) -> dict[int, dict]:
	"""Build a frame_index → hand keypoints lookup dict.

	Returns dict like::

	{0: {"h0": [[x,y]21], "h1": [[x,y]21]}, 30: {"h0": [...], "h1": None}, ...}
	"""
	df = pd.read_parquet(parquet_path)
	lookup: dict[int, dict] = {}

	has_h0 = "hand_0_keypoints" in df.columns
	has_h1 = "hand_1_keypoints" in df.columns

	if not has_h0 and not has_h1:
	return lookup

	for _, row in df.iterrows():
	fi = int(row["frame_index"])
	entry: dict = {}
	for hk, col in (("h0", "hand_0_keypoints"), ("h1", "hand_1_keypoints")):
	if col not in df.columns:
	continue
	val = row[col]
	if val is None or (isinstance(val, float) and pd.isna(val)):
	entry[hk] = None
	continue
	kp_list = val.tolist() if hasattr(val, "tolist") else val
	if isinstance(kp_list, list) and len(kp_list) >= 21:
	entry[hk] = [[float(p[0]), float(p[1])] for p in kp_list[:21]]
	else:
	entry[hk] = None
	lookup[fi] = entry

	return lookup


	# ═══════════════════════════════════════════════════════════════
	# Drawing
	# ═══════════════════════════════════════════════════════════════

	def _get_drawing_func():
	"""Return the best available drawing function.

	Tries Google MediaPipe official drawing_utils first,
	falls back to OpenCV-based rendering.
	"""
	try:
	from mediapipe.tasks.python.vision import drawing_utils, drawing_styles
	from mediapipe.tasks.python.vision.hand_landmarker import HandLandmarksConnections
	from mediapipe.tasks.python.components.containers import landmark as lm

	def _draw_mediapipe(frame_bgr, kp_lookup, frame_index):
	"""Draw hand skeletons on a BGR frame using Google MediaPipe official API."""
	entry = kp_lookup.get(frame_index)
	if not entry:
	return

	W, H = 640, 480 # source resolution

	for hk in ("h0", "h1"):
	kp = entry.get(hk)
	if not kp or len(kp) < 21:
	continue

	landmarks = []
	for pt in kp[:21]:
	l = lm.NormalizedLandmark()
	l.x = pt[0] / W
	l.y = pt[1] / H
	l.z = 0.0
	landmarks.append(l)

	drawing_utils.draw_landmarks(
	frame_bgr,
	landmarks,
	HandLandmarksConnections.HAND_CONNECTIONS,
	drawing_styles.get_default_hand_landmarks_style(),
	drawing_styles.get_default_hand_connections_style(),
	)

	return _draw_mediapipe, "mediapipe"

	except ImportError:
	pass

	# ── OpenCV fallback ──
	def _draw_opencv(frame_bgr, kp_lookup, frame_index):
	"""Draw hand skeletons using OpenCV (near-MediaPipe replica)."""
	entry = kp_lookup.get(frame_index)
	if not entry:
	return

	import cv2

	connections = [
	(0,1),(1,2),(2,3),(3,4),(0,5),(5,6),(6,7),(7,8),
	(0,9),(9,10),(10,11),(11,12),(0,13),(13,14),(14,15),(15,16),
	(0,17),(17,18),(18,19),(19,20),(5,9),(9,13),(13,17),
	]
	# BGR colours per finger (matches MediaPipe style)
	finger_colors = [
	(255, 48, 48), # thumb → red
	(128, 64, 128), # index → purple
	(48, 255, 255), # middle → yellow
	(48, 255, 48), # ring → green
	(48, 208, 255), # pinky → orange
	]

	for hk in ("h0", "h1"):
	kp = entry.get(hk)
	if not kp or len(kp) < 21:
	continue
	pts = [(int(p[0]), int(p[1])) for p in kp[:21]]

	# Palm connections (gray)
	for i, j in [(0, 1), (0, 17), (1, 5), (5, 9), (9, 13), (13, 17)]:
	cv2.line(frame_bgr, pts[i], pts[j], (128, 128, 128), 3, cv2.LINE_AA)

	# Finger connections (colored)
	for f in range(5):
	c = finger_colors[f]
	base = f * 4 + 1
	for i in range(3):
	cv2.line(frame_bgr, pts[base + i], pts[base + i + 1],
	c, 2, cv2.LINE_AA)

	# Landmark dots: white border + colored fill
	radii = [7, 4, 4, 3, 2, 4, 4, 3, 2, 4, 4, 3, 2,
	4, 4, 3, 2, 4, 4, 3, 2]
	for i, (x, y) in enumerate(pts):
	if i == 0:
	fc = (255, 255, 255) # wrist white
	else:
	fc = finger_colors[(i - 1) // 4]
	r = radii[i]
	cv2.circle(frame_bgr, (x, y), r + 1, (255, 255, 255), -1, cv2.LINE_AA)
	cv2.circle(frame_bgr, (x, y), max(1, r - 1), fc, -1, cv2.LINE_AA)

	return _draw_opencv, "opencv"


	# ═══════════════════════════════════════════════════════════════
	# Main entry point
	# ═══════════════════════════════════════════════════════════════

	async def render_skeleton_video(
	video_path: Path,
	session_dir: Path,
	) -> Path \| None:
	"""Render hand skeletons onto a video, returning the output path.

	Args:
	video_path: Path to the source MP4 file.
	session_dir: Session root directory (contains data/ with parquet).

	Returns:
	Path to the skeleton video, or ``None`` if no hand data exists.
	"""
	import cv2

	# 1. Find and validate hand data
	parquet_path = _find_data_parquet(session_dir)
	if parquet_path is None:
	logger.info("No data parquet found in %s — skipping skeleton render", session_dir)
	return None

	if not _has_hand_keypoints(parquet_path):
	logger.info("No hand keypoint columns in %s — skipping skeleton render", parquet_path.name)
	return None

	# 2. Build keypoint lookup (frame_index → keypoints)
	kp_lookup = _build_keypoint_lookup(parquet_path)
	keypoint_frames = sum(1 for v in kp_lookup.values() if v.get("h0") or v.get("h1"))
	if keypoint_frames == 0:
	logger.info("No frames with hand keypoints — skipping skeleton render")
	return None

	logger.info(
	"Rendering skeleton video: %d / %d total rows have hand keypoints",
	keypoint_frames, len(kp_lookup),
	)

	# 3. Get drawing function
	draw_func, backend = _get_drawing_func()
	logger.info("Using drawing backend: %s", backend)

	# 4. Open video
	cap = cv2.VideoCapture(str(video_path))
	if not cap.isOpened():
	logger.error("Cannot open video: %s", video_path)
	return None

	fps = cap.get(cv2.CAP_PROP_FPS)
	if fps <= 0:
	fps = 30.0
	width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
	height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
	total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

	logger.info("Video: %dx%d @ %.1f fps, %d frames", width, height, fps, total_frames)

	# 5. Output path: {stem}_skeleton{ext}
	output_path = video_path.parent / f"{video_path.stem}_skeleton{video_path.suffix}"

	# 6. FFmpeg H.264 encoder (browser-compatible)
	ffmpeg_cmd = [
	"ffmpeg", "-y",
	"-f", "rawvideo", "-vcodec", "rawvideo",
	"-s", f"{width}x{height}", "-pix_fmt", "bgr24", "-r", str(fps),
	"-i", "pipe:0",
	"-c:v", "libx264", "-preset", "fast", "-crf", "23",
	"-pix_fmt", "yuv420p",
	str(output_path),
	]
	proc = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE, stderr=subprocess.DEVNULL)
	frame_size = width * height * 3 # BGR24 = 3 bytes per pixel

	# 7. Process frames
	frame_idx = 0
	rendered_frames = 0
	while True:
	ret, frame = cap.read()
	if not ret:
	break

	draw_func(frame, kp_lookup, frame_idx)
	proc.stdin.write(frame.tobytes())

	if kp_lookup.get(frame_idx):
	rendered_frames += 1

	frame_idx += 1
	if frame_idx % 50 == 0:
	logger.debug("Skeleton render: %d / %d frames", frame_idx, total_frames)

	cap.release()
	proc.stdin.close()
	proc.wait()

	logger.info(
	"Skeleton video written: %s (%d frames rendered, %d total)",
	output_path.name, rendered_frames, frame_idx,
	)
	return output_path