Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-09-15 14:13:49 +08:00
commit 7e4ef6f98b
7002 changed files with 1452249 additions and 0 deletions
@@ -0,0 +1,124 @@
+import os
+import joblib
+import json
+
+import cv2
+import numpy as np
+
+
+def export_hamer_predictions(res_path, target_dir):
+    """
+    exports hamer output results into a directory for each track
+    of JSON files with each frame's mano prediction
+    """
+    os.makedirs(target_dir, exist_ok=True)
+    print(f"exporting hamer predictions from {res_path} to {target_dir}")
+
+    tracklet_data = joblib.load(res_path)
+    for frame, track_data in tracklet_data.items():
+        name = os.path.splitext(frame.split("/")[-1])[0]
+        print(track_data['tid'])
+        track_dicts = unpack_frame(track_data)
+        for tid, pred_dict in track_dicts.items():
+            track_dir = os.path.join(target_dir, f"{tid:03d}")
+            os.makedirs(track_dir, exist_ok=True)
+            pred_path = os.path.join(track_dir, f"{name}_mano.json")
+            with open(pred_path, "w") as f:
+                json.dump(pred_dict, f, indent=1)
+
+
+def export_vitpose_keypoints(res_path, target_dir):
+    """
+    exports each track of each frame into its own json of keypoints
+    """
+    os.makedirs(target_dir, exist_ok=True)
+    print(f"exporting keypoints from {res_path} to {target_dir}")
+
+    tracklet_data = joblib.load(res_path)
+    for frame, track_data in tracklet_data.items():
+        name = os.path.splitext(frame.split("/")[-1])[0]
+        tids = track_data["tid"]
+        valid_tids = track_data["tracked_ids"]
+        kps_all = track_data["extra_data"]
+        # mask_paths = track_data["mask_name"]
+        for i, tid in enumerate(tids):
+            if tid not in valid_tids:
+                continue
+            track_dir = os.path.join(target_dir, f"{tid:03d}")
+            os.makedirs(track_dir, exist_ok=True)
+            kp_path = os.path.join(track_dir, f"{name}_keypoints.json")
+            kp_dict = {
+                "people": [
+                    {
+                        "pose_keypoints_2d": kps_all[i],#.tolist(),
+                        # "mask_path": mask_paths[i],
+                    }
+                ],
+            }
+            with open(kp_path, "w") as f:
+                json.dump(kp_dict, f, indent=1)
+
+
+def export_shot_changes(res_path, target_path):
+    """
+    Exports the phalp output of shot changes into a JSON file
+    dict with the shot index for each frame in the sequence
+    """
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+    print(f"exporting shot changes to {target_path}")
+
+    tracklet_data = joblib.load(res_path)
+    frames = sorted(tracklet_data.keys())
+    # print(frames)
+    # print(tracklet_data[frames[0]].keys())
+    # for frame in frames:
+    #     print(frame, tracklet_data[frame]["shot"])
+    shots = np.cumsum([tracklet_data[frame]["shot"] for frame in frames])
+    shots = shots.astype(int).tolist()
+    shot_dict = {frame.split("/")[-1]: shot for frame, shot in zip(frames, shots)}
+    with open(target_path, "w") as f:
+        json.dump(shot_dict, f, indent=1)
+
+
+def unpack_frame(track_data):
+    """
+    Unpack one frame of track data from the phalp output into JSON files for each track.
+    Track data contains body pose, global translation and rotation,
+    and keypoints for all tracks in the frame.
+    Each track is unpacked into folders with a JSON for each frame.
+    """
+    out_dict = {}
+    for i, tid in enumerate(track_data["tid"]):
+        if tid not in track_data["tracked_ids"]:
+            continue
+
+        cam_trans = track_data["cam_trans"][i].squeeze()  # (3,)
+
+        mano = track_data["mano"][i]
+        betas = mano["betas"].squeeze()  # (10,)
+        hand_pose = mano["hand_pose"].squeeze()  # (23, 3, 3)
+        is_right = mano["is_right"]
+        print(is_right, tid, track_data["tid"], track_data["tracked_ids"])
+        assert is_right == tid, f'{is_right}, {tid}'
+        hand_pose_aa = np.stack(
+            [cv2.Rodrigues(x)[0].squeeze() for x in hand_pose], axis=0
+        )  # (23, 3)
+        global_orient = mano["global_orient"].squeeze()  # (3, 3)
+        global_orient_aa = cv2.Rodrigues(global_orient)[0].squeeze()  # (3,)
+
+        out_dict[tid] = {
+            "betas": betas.tolist(),
+            "body_pose": hand_pose_aa.tolist(),
+            "global_orient": global_orient_aa.tolist(),
+            "cam_trans": cam_trans.tolist(),
+            "is_right": is_right
+        }
+
+    return out_dict
+
+
+def export_sequence_results(hamer_res_path, track_dir, shot_path):
+    print('exporting sequence results in export_hamer.py...')
+    export_hamer_predictions(hamer_res_path, track_dir)
+    export_vitpose_keypoints(hamer_res_path, track_dir)
+    export_shot_changes(hamer_res_path, shot_path)