#!/usr/bin/env python3 """ 手部关键点提取 CLI —— 对已录制会话离线处理。 用法: python scripts/process_hands.py kpts /path/to/session --mode bare python scripts/process_hands.py label /path/to/session python scripts/process_hands.py all /path/to/session --mode bare """ import argparse import os import sys import subprocess import shutil import numpy as np import cv2 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from core.hand_tracking import process_session, label_session, load_hand_kpts, draw_kpts_overlay def _progress_bar(cur: int, total: int, width: int = 40): ratio = min(cur / max(total, 1), 1.0) filled = int(width * ratio) bar = "█" * filled + "░" * (width - filled) sys.stderr.write(f"\r [{bar}] {cur}/{total} ({ratio*100:.0f}%)") sys.stderr.flush() def cmd_kpts(args): """提取手部关键点（2D + 可选 3D）。""" print(f"提取手部关键点: {args.session}") print(f" 模式: {args.mode}, 设备: {args.device}") sys.stderr.write("\n") result = process_session( session_path=args.session, mode=args.mode, det_device=args.device, pose_device=args.device, progress_cb=_progress_bar, status_cb=lambda msg: print(f" [{msg}]"), ) sys.stderr.write("\n") if result["success"]: print(f"✅ 完成: {result['frames']} 帧, " f"{result['elapsed']:.1f}s, {result['fps']:.1f} fps") else: print(f"❌ 失败: {result['error']}") sys.exit(1) def cmd_label(args): """手势自动标注。""" print(f"手势标注: {args.session}") result = label_session( session_path=args.session, progress_cb=_progress_bar, ) sys.stderr.write("\n") if result["success"]: print(f"✅ 完成: {result['frames']} 帧, {result['elapsed']:.1f}s") else: print(f"❌ 失败: {result['error']}") sys.exit(1) def cmd_all(args): """先提取关键点，再标注。""" cmd_kpts(args) print() cmd_label(args) def cmd_show(args): """显示已有关键点信息。""" kpts = load_hand_kpts(args.session) if kpts: n = len(kpts) sample = kpts[min(kpts.keys())] print(f"手部关键点数据: {n} 帧") print(f" 维度: {len(sample['hand_data'])}") print(f" 最多手数: {sample['num_hands']}") else: print("❌ 未找到手部关键点数据") sys.exit(1) def _find_video(session_path: str) -> str: """在 session 目录下找第一个可用的 RGB 视频文件。""" from core.helpers import detect_session_format, egodata_video_path, egodata_metadata_path import json vdir = os.path.join(session_path, "videos") if not os.path.isdir(vdir): return "" fmt = detect_session_format(session_path) # 先尝试从 metadata 获取相机列表 if fmt == "egodata": meta_path = egodata_metadata_path(session_path) if os.path.isfile(meta_path): with open(meta_path) as f: cameras = json.load(f).get("cameras", {}) for cam_name in cameras: if cameras[cam_name].get("type") == "depth": continue for vp in [egodata_video_path(session_path, cam_name), os.path.join(vdir, f"{cam_name}.mp4")]: if os.path.isfile(vp): return vp # 直接扫描 videos/ 目录 for entry in sorted(os.listdir(vdir)): if "depth" in entry.lower(): continue full = os.path.join(vdir, entry) if os.path.isfile(full) and entry.endswith(".mp4"): return full # 子目录结构 sub_dir = os.path.join(full, "chunk-0000") if os.path.isdir(sub_dir): for fname in os.listdir(sub_dir): if fname.endswith(".mp4"): return os.path.join(sub_dir, fname) return "" def cmd_render(args): """将手部关键点叠加到视频上，导出可视化视频。""" import time from core.helpers import keypoints_video_dir session_path = args.session kpts = load_hand_kpts(session_path) if not kpts: print("❌ 未找到手部关键点数据，请先运行 kpts 命令") sys.exit(1) video_path = _find_video(session_path) if not video_path: print("❌ 未找到可播放的视频文件") sys.exit(1) cam_name = os.path.splitext(os.path.basename(video_path))[0] out_dir = keypoints_video_dir(session_path) os.makedirs(out_dir, exist_ok=True) out_path = os.path.join(out_dir, f"{cam_name}_hand_kpts.mp4") cap = cv2.VideoCapture(video_path) total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) fps = cap.get(cv2.CAP_PROP_FPS) if fps <= 0: fps = 30.0 ret, first = cap.read() if not ret: print("❌ 无法读取视频帧") sys.exit(1) h, w = first.shape[:2] cap.set(cv2.CAP_PROP_POS_FRAMES, 0) # ffmpeg 优先，cv2 回退 writer = None ffmpeg_proc = None if shutil.which("ffmpeg"): ffmpeg_cmd = [ "ffmpeg", "-y", "-f", "rawvideo", "-vcodec", "rawvideo", "-s", f"{w}x{h}", "-pix_fmt", "bgr24", "-r", str(fps), "-i", "-", "-c:v", "libx264", "-preset", "fast", "-crf", "23", "-pix_fmt", "yuv420p", out_path, ] try: ffmpeg_proc = subprocess.Popen( ffmpeg_cmd, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) except Exception: ffmpeg_proc = None if not ffmpeg_proc: for codec in ["avc1", "mp4v"]: fourcc = cv2.VideoWriter_fourcc(*codec) writer = cv2.VideoWriter(out_path, fourcc, fps, (w, h)) if writer.isOpened(): break print(f"输入视频: {video_path} ({total} 帧, {fps} fps)") print(f"手部关键点: {len(kpts)} 帧") print(f"输出: {out_path}") t0 = time.perf_counter() fi = 0 while True: ret, frame = cap.read() if not ret: break kpt = kpts.get(fi) if kpt is not None and np.any(kpt["hand_data"] > 0): frame = draw_kpts_overlay(frame, kpt["hand_data"], kpt["track_ids"]) if ffmpeg_proc: try: ffmpeg_proc.stdin.write(frame.tobytes()) except BrokenPipeError: break else: writer.write(frame) fi += 1 if fi % 50 == 0: _progress_bar(fi, total) cap.release() if ffmpeg_proc: ffmpeg_proc.stdin.close() ffmpeg_proc.wait() else: writer.release() elapsed = time.perf_counter() - t0 _progress_bar(total, total) sys.stderr.write("\n") size_mb = os.path.getsize(out_path) / (1024 * 1024) print(f"✅ 完成: {out_path}") print(f" {fi} 帧, {elapsed:.1f}s ({fi / elapsed:.1f} fps), {size_mb:.1f} MB") def main(): parser = argparse.ArgumentParser( description="手部关键点提取与标注", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例: python scripts/process_hands.py kpts data/recordings/Test005/episode_000001 --mode bare python scripts/process_hands.py render data/recordings/Test005/episode_000001 python scripts/process_hands.py all data/recordings/Test005/episode_000001 --mode bare python scripts/process_hands.py show data/recordings/Test005/episode_000001 """, ) sub = parser.add_subparsers(dest="command", help="子命令") p_kpts = sub.add_parser("kpts", help="提取手部关键点") p_kpts.add_argument("session", help="会话目录路径") p_kpts.add_argument("--mode", default="glove", choices=["glove", "bare"], help="追踪模式") p_kpts.add_argument("--device", default="cuda", help="推理设备") p_label = sub.add_parser("label", help="手势自动标注") p_label.add_argument("session", help="会话目录路径") p_all = sub.add_parser("all", help="提取 + 标注") p_all.add_argument("session", help="会话目录路径") p_all.add_argument("--mode", default="glove", choices=["glove", "bare"], help="追踪模式") p_all.add_argument("--device", default="cuda", help="推理设备") p_show = sub.add_parser("show", help="查看已有数据") p_show.add_argument("session", help="会话目录路径") p_render = sub.add_parser("render", help="导出关键点可视化视频") p_render.add_argument("session", help="会话目录路径") args = parser.parse_args() if args.command == "kpts": cmd_kpts(args) elif args.command == "label": cmd_label(args) elif args.command == "all": cmd_all(args) elif args.command == "show": cmd_show(args) elif args.command == "render": cmd_render(args) else: parser.print_help() sys.exit(1) if __name__ == "__main__": main()