Video2MC / data /data_utils.py
Sapphire-356's picture
add: Video2MC
95f8bbc
raw
history blame
2.81 kB
# Copyright (c) 2018-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import h5py
import numpy as np
mpii_metadata = {
'layout_name': 'mpii',
'num_joints': 16,
'keypoints_symmetry': [
[3, 4, 5, 13, 14, 15],
[0, 1, 2, 10, 11, 12],
]
}
coco_metadata = {
'layout_name': 'coco',
'num_joints': 17,
'keypoints_symmetry': [
[1, 3, 5, 7, 9, 11, 13, 15],
[2, 4, 6, 8, 10, 12, 14, 16],
]
}
h36m_metadata = {
'layout_name': 'h36m',
'num_joints': 17,
'keypoints_symmetry': [
[4, 5, 6, 11, 12, 13],
[1, 2, 3, 14, 15, 16],
]
}
humaneva15_metadata = {
'layout_name': 'humaneva15',
'num_joints': 15,
'keypoints_symmetry': [
[2, 3, 4, 8, 9, 10],
[5, 6, 7, 11, 12, 13]
]
}
humaneva20_metadata = {
'layout_name': 'humaneva20',
'num_joints': 20,
'keypoints_symmetry': [
[3, 4, 5, 6, 11, 12, 13, 14],
[7, 8, 9, 10, 15, 16, 17, 18]
]
}
def suggest_metadata(name):
names = []
for metadata in [mpii_metadata, coco_metadata, h36m_metadata, humaneva15_metadata, humaneva20_metadata]:
if metadata['layout_name'] in name:
return metadata
names.append(metadata['layout_name'])
raise KeyError('Cannot infer keypoint layout from name "{}". Tried {}.'.format(name, names))
def import_detectron_poses(path):
# Latin1 encoding because Detectron runs on Python 2.7
data = np.load(path, encoding='latin1')
kp = data['keypoints']
bb = data['boxes']
results = []
for i in range(len(bb)):
if len(bb[i][1]) == 0:
assert i > 0
# Use last pose in case of detection failure
results.append(results[-1])
continue
best_match = np.argmax(bb[i][1][:, 4])
# import ipdb;ipdb.set_trace()
keypoints = kp[i][1][best_match].T.copy()
results.append(keypoints)
results = np.array(results)
# return results[:, :, 4:6] # Soft-argmax
return results[:, :, [0, 1, 3]] # Argmax + score
def my_pose(path):
data = np.load(path, encoding='latin1')
def import_cpn_poses(path):
data = np.load(path)
kp = data['keypoints']
return kp[:, :, :2]
def import_sh_poses(path):
with h5py.File(path) as hf:
positions = hf['poses'].value
return positions.astype('float32')
def suggest_pose_importer(name):
if 'detectron' in name:
return import_detectron_poses
if 'cpn' in name:
return import_cpn_poses
if 'sh' in name:
return import_sh_poses
raise KeyError('Cannot infer keypoint format from name "{}". Tried detectron, cpn, sh.'.format(name))