125 lines
4.5 KiB
Python
125 lines
4.5 KiB
Python
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)
|