68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
|
import json
|
||
|
|
import os
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
def load_rtmpose_json(filepath):
|
||
|
|
with open(filepath, 'r') as f:
|
||
|
|
data = json.load(f)
|
||
|
|
return data
|
||
|
|
|
||
|
|
def convert_to_pose2sim_format(frame_data):
|
||
|
|
pose2sim_frame = {
|
||
|
|
"version": 1.3,
|
||
|
|
"people": []
|
||
|
|
}
|
||
|
|
|
||
|
|
for instance in frame_data['instances']:
|
||
|
|
keypoints = instance['keypoints']
|
||
|
|
scores = instance['keypoint_scores']
|
||
|
|
|
||
|
|
pose_keypoints_2d = []
|
||
|
|
for idx, point in enumerate(keypoints):
|
||
|
|
pose_keypoints_2d.extend(point)
|
||
|
|
pose_keypoints_2d.append(scores[idx])
|
||
|
|
|
||
|
|
person_data = {
|
||
|
|
"person_id": [-1],
|
||
|
|
"pose_keypoints_2d": pose_keypoints_2d,
|
||
|
|
"face_keypoints_2d": [],
|
||
|
|
"hand_left_keypoints_2d": [],
|
||
|
|
"hand_right_keypoints_2d": [],
|
||
|
|
"pose_keypoints_3d": [],
|
||
|
|
"face_keypoints_3d": [],
|
||
|
|
"hand_left_keypoints_3d": [],
|
||
|
|
"hand_right_keypoints_3d": []
|
||
|
|
}
|
||
|
|
|
||
|
|
pose2sim_frame["people"].append(person_data)
|
||
|
|
|
||
|
|
return pose2sim_frame
|
||
|
|
|
||
|
|
def save_pose2sim_json(data, frame_id, output_dir):
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
output_filepath = os.path.join(output_dir, f"frame_{frame_id}.json")
|
||
|
|
with open(output_filepath, 'w') as f:
|
||
|
|
json.dump(data, f, indent=4)
|
||
|
|
|
||
|
|
def convert_rtmpose_to_pose2sim(input_filepath, output_dir):
|
||
|
|
rtmpose_data = load_rtmpose_json(input_filepath)
|
||
|
|
frames = rtmpose_data['instance_info']
|
||
|
|
for frame in frames:
|
||
|
|
frame_id = frame['frame_id']
|
||
|
|
pose2sim_frame = convert_to_pose2sim_format(frame)
|
||
|
|
save_pose2sim_json(pose2sim_frame, frame_id, output_dir)
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description='Convert RTMPose JSON to Pose2Sim JSON format.')
|
||
|
|
parser.add_argument('-i', '--input', required=True, help='Path to the input RTMPose JSON file.')
|
||
|
|
parser.add_argument('-o', '--output', required=True, help='Path to the output directory where Pose2Sim JSON files will be saved.')
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
input_filepath = args.input
|
||
|
|
output_dir = args.output
|
||
|
|
|
||
|
|
convert_rtmpose_to_pose2sim(input_filepath, output_dir)
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|