# -----------------------------------------------------------------------------
# Adapted from https://github.com/anibali/h36m-fetch
# Original license: Copyright (c) Aiden Nibali, under the Apache License.
# -----------------------------------------------------------------------------
import argparse
import os
import pickle
import tarfile
import xml.etree.ElementTree as ET
from os.path import join
import cv2
import numpy as np
from spacepy import pycdf
class PreprocessH36m:
"""Preprocess Human3.6M dataset.
Args:
metadata (str): Path to metadata.xml.
original_dir (str): Directory of the original dataset with all files
compressed. Specifically, .tgz files belonging to subject 1
should be placed under the subdirectory 's1'.
extracted_dir (str): Directory of the extracted files. If not given, it
will be placed under the same parent directory as original_dir.
processed_der (str): Directory of the processed files. If not given, it
will be placed under the same parent directory as original_dir.
sample_rate (int): Downsample FPS to `1 / sample_rate`. Default: 5.
"""
def __init__(self,
metadata,
original_dir,
extracted_dir=None,
processed_dir=None,
sample_rate=5):
self.metadata = metadata
self.original_dir = original_dir
self.sample_rate = sample_rate
if extracted_dir is None:
self.extracted_dir = join(
os.path.dirname(os.path.abspath(self.original_dir)),
'extracted')
else:
self.extracted_dir = extracted_dir
if processed_dir is None:
self.processed_dir = join(
os.path.dirname(os.path.abspath(self.original_dir)),
'processed')
else:
self.processed_dir = processed_dir
self.subjects = []
self.sequence_mappings = {}
self.action_names = {}
self.camera_ids = []
self._load_metadata()
self.subjects_annot = ['S1', 'S5', 'S6', 'S7', 'S8', 'S9', 'S11']
self.subjects_splits = {
'train': ['S1', 'S5', 'S6', 'S7', 'S8'],
'test': ['S9', 'S11']
}
self.extract_files = ['Videos', 'D2_Positions', 'D3_Positions_mono']
self.movable_joints = [
0, 1, 2, 3, 6, 7, 8, 12, 13, 14, 15, 17, 18, 19, 25, 26, 27
]
self.scale_factor = 1.2
self.image_sizes = {
'54138969': {
'width': 1000,
'height': 1002
},
'55011271': {
'width': 1000,
'height': 1000
},
'58860488': {
'width': 1000,
'height': 1000
},
'60457274': {
'width': 1000,
'height': 1002
}
}
def extract_tgz(self):
"""Extract files from self.extrct_files."""
os.makedirs(self.extracted_dir, exist_ok=True)
for subject in self.subjects_annot:
cur_dir = join(self.original_dir, subject.lower())
for file in self.extract_files:
filename = join(cur_dir, file + '.tgz')
print(f'Extracting {filename} ...')
with tarfile.open(filename) as tar:
tar.extractall(self.extracted_dir)
print('Extraction done.\n')
def generate_cameras_file(self):
"""Generate cameras.pkl which contains camera parameters for 11
subjects each with 4 cameras."""
cameras = {}
for subject in range(1, 12):
for camera in range(4):
key = (f'S{subject}', self.camera_ids[camera])
cameras[key] = self._get_camera_params(camera, subject)
out_file = join(self.processed_dir, 'annotation_body3d', 'cameras.pkl')
with open(out_file, 'wb') as fout:
pickle.dump(cameras, fout)
print(f'Camera parameters have been written to "{out_file}".\n')
def generate_annotations(self):
"""Generate annotations for training and testing data."""
output_dir = join(self.processed_dir, 'annotation_body3d',
f'fps{50 // self.sample_rate}')
os.makedirs(output_dir, exist_ok=True)
for data_split in ('train', 'test'):
imgnames_all = []
centers_all = []
scales_all = []
kps2d_all = []
kps3d_all = []
for subject in self.subjects_splits[data_split]:
for action, subaction in self.sequence_mappings[subject].keys(
):
if action == '1':
# exclude action "_ALL"
continue
for camera in self.camera_ids:
imgnames, centers, scales, kps2d, kps3d\
= self._load_annotations(
subject, action, subaction, camera)
imgnames_all.append(imgnames)
centers_all.append(centers)
scales_all.append(scales)
kps2d_all.append(kps2d)
kps3d_all.append(kps3d)
imgnames_all = np.concatenate(imgnames_all)
centers_all = np.concatenate(centers_all)
scales_all = np.concatenate(scales_all)
kps2d_all = np.concatenate(kps2d_all)
kps3d_all = np.concatenate(kps3d_all)
out_file = join(output_dir, f'h36m_{data_split}.npz')
np.savez(
out_file,
imgname=imgnames_all,
center=centers_all,
scale=scales_all,
part=kps2d_all,
S=kps3d_all)
print(
f'All annotations of {data_split}ing data have been written to'
f' "{out_file}". {len(imgnames_all)} samples in total.\n')
if data_split == 'train':
kps_3d_all = kps3d_all[..., :3] # remove visibility
mean_3d, std_3d = self._get_pose_stats(kps_3d_all)
kps_2d_all = kps2d_all[..., :2] # remove visibility
mean_2d, std_2d = self._get_pose_stats(kps_2d_all)
# centered around root
# the root keypoint is 0-index
kps_3d_rel = kps_3d_all[..., 1:, :] - kps_3d_all[..., :1, :]
mean_3d_rel, std_3d_rel = self._get_pose_stats(kps_3d_rel)
kps_2d_rel = kps_2d_all[..., 1:, :] - kps_2d_all[..., :1, :]
mean_2d_rel, std_2d_rel = self._get_pose_stats(kps_2d_rel)
stats = {
'joint3d_stats': {
'mean': mean_3d,
'std': std_3d
},
'joint2d_stats': {
'mean': mean_2d,
'std': std_2d
},
'joint3d_rel_stats': {
'mean': mean_3d_rel,
'std': std_3d_rel
},
'joint2d_rel_stats': {
'mean': mean_2d_rel,
'std': std_2d_rel
}
}
for name, stat_dict in stats.items():
out_file = join(output_dir, f'{name}.pkl')
with open(out_file, 'wb') as f:
pickle.dump(stat_dict, f)
print(f'Create statistic data file: {out_file}')
@staticmethod
def _get_pose_stats(kps):
"""Get statistic information `mean` and `std` of pose data.
Args:
kps (ndarray): keypoints in shape [..., K, C] where K and C is
the keypoint category number and dimension.
Returns:
mean (ndarray): [K, C]
"""
assert kps.ndim > 2
K, C = kps.shape[-2:]
kps = kps.reshape(-1, K, C)
mean = kps.mean(axis=0)
std = kps.std(axis=0)
return mean, std
def _load_metadata(self):
"""Load meta data from metadata.xml."""
assert os.path.exists(self.metadata)
tree = ET.parse(self.metadata)
root = tree.getroot()
for i, tr in enumerate(root.find('mapping')):
if i == 0:
_, _, *self.subjects = [td.text for td in tr]
self.sequence_mappings \
= {subject: {} for subject in self.subjects}
elif i < 33:
action_id, subaction_id, *prefixes = [td.text for td in tr]
for subject, prefix in zip(self.subjects, prefixes):
self.sequence_mappings[subject][(action_id, subaction_id)]\
= prefix
for i, elem in enumerate(root.find('actionnames')):
action_id = str(i + 1)
self.action_names[action_id] = elem.text
self.camera_ids \
= [elem.text for elem in root.find('dbcameras/index2id')]
w0 = root.find('w0')
self.cameras_raw = [float(num) for num in w0.text[1:-1].split()]
def _get_base_filename(self, subject, action, subaction, camera):
"""Get base filename given subject, action, subaction and camera."""
return f'{self.sequence_mappings[subject][(action, subaction)]}' + \
f'.{camera}'
def _get_camera_params(self, camera, subject):
"""Get camera parameters given camera id and subject id."""
metadata_slice = np.zeros(15)
start = 6 * (camera * 11 + (subject - 1))
metadata_slice[:6] = self.cameras_raw[start:start + 6]
metadata_slice[6:] = self.cameras_raw[265 + camera * 9 - 1:265 +
(camera + 1) * 9 - 1]
# extrinsics
x, y, z = -metadata_slice[0], metadata_slice[1], -metadata_slice[2]
R_x = np.array([[1, 0, 0], [0, np.cos(x), np.sin(x)],
[0, -np.sin(x), np.cos(x)]])
R_y = np.array([[np.cos(y), 0, np.sin(y)], [0, 1, 0],
[-np.sin(y), 0, np.cos(y)]])
R_z = np.array([[np.cos(z), np.sin(z), 0], [-np.sin(z),
np.cos(z), 0], [0, 0, 1]])
R = (R_x @ R_y @ R_z).T
T = metadata_slice[3:6].reshape(-1, 1)
# convert unit from millimeter to meter
T *= 0.001
# intrinsics
c = metadata_slice[8:10, None]
f = metadata_slice[6:8, None]
# distortion
k = metadata_slice[10:13, None]
p = metadata_slice[13:15, None]
return {
'R': R,
'T': T,
'c': c,
'f': f,
'k': k,
'p': p,
'w': self.image_sizes[self.camera_ids[camera]]['width'],
'h': self.image_sizes[self.camera_ids[camera]]['height'],
'name': f'camera{camera + 1}',
'id': self.camera_ids[camera]
}
def _load_annotations(self, subject, action, subaction, camera):
"""Load annotations for a sequence."""
subj_dir = join(self.extracted_dir, subject)
basename = self._get_base_filename(subject, action, subaction, camera)
# load 2D keypoints
with pycdf.CDF(
join(subj_dir, 'MyPoseFeatures', 'D2_Positions',
basename + '.cdf')) as cdf:
kps_2d = np.array(cdf['Pose'])
num_frames = kps_2d.shape[1]
kps_2d = kps_2d.reshape((num_frames, 32, 2))[::self.sample_rate,
self.movable_joints]
kps_2d = np.concatenate([kps_2d, np.ones((len(kps_2d), 17, 1))],
axis=2)
# load 3D keypoints
with pycdf.CDF(
join(subj_dir, 'MyPoseFeatures', 'D3_Positions_mono',
basename + '.cdf')) as cdf:
kps_3d = np.array(cdf['Pose'])
kps_3d = kps_3d.reshape(
(num_frames, 32, 3))[::self.sample_rate,
self.movable_joints] / 1000.
kps_3d = np.concatenate([kps_3d, np.ones((len(kps_3d), 17, 1))],
axis=2)
# calculate bounding boxes
bboxes = np.stack([
np.min(kps_2d[:, :, 0], axis=1),
np.min(kps_2d[:, :, 1], axis=1),
np.max(kps_2d[:, :, 0], axis=1),
np.max(kps_2d[:, :, 1], axis=1)
],
axis=1)
centers = np.stack([(bboxes[:, 0] + bboxes[:, 2]) / 2,
(bboxes[:, 1] + bboxes[:, 3]) / 2],
axis=1)
scales = self.scale_factor * np.max(
bboxes[:, 2:] - bboxes[:, :2], axis=1) / 200
# extract frames and save imgnames
imgnames = []
video_path = join(subj_dir, 'Videos', basename + '.mp4')
sub_base = subject + '_' + basename.replace(' ', '_')
img_dir = join(self.processed_dir, 'images', subject, sub_base)
os.makedirs(img_dir, exist_ok=True)
prefix = join(subject, sub_base, sub_base)
cap = cv2.VideoCapture(video_path)
i = 0
while True:
success, img = cap.read()
if not success:
break
if i % self.sample_rate == 0:
imgname = f'{prefix}_{i + 1:06d}.jpg'
imgnames.append(imgname)
dest_path = join(self.processed_dir, 'images', imgname)
if not os.path.exists(dest_path):
cv2.imwrite(dest_path, img)
if len(imgnames) == len(centers):
break
i += 1
cap.release()
imgnames = np.array(imgnames)
print(f'Annoatations for sequence "{subject} {basename}" are loaded. '
f'{len(imgnames)} samples in total.')
return imgnames, centers, scales, kps_2d, kps_3d
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
'--metadata', type=str, required=True, help='Path to metadata.xml')
parser.add_argument(
'--original',
type=str,
required=True,
help='Directory of the original dataset with all files compressed. '
'Specifically, .tgz files belonging to subject 1 should be placed '
'under the subdirectory \"s1\".')
parser.add_argument(
'--extracted',
type=str,
default=None,
help='Directory of the extracted files. If not given, it will be '
'placed under the same parent directory as original_dir.')
parser.add_argument(
'--processed',
type=str,
default=None,
help='Directory of the processed files. If not given, it will be '
'placed under the same parent directory as original_dir.')
parser.add_argument(
'--sample_rate',
type=int,
default=5,
help='Downsample FPS to `1 / sample_rate`. Default: 5.')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
h36m = PreprocessH36m(
metadata=args.metadata,
original_dir=args.original,
extracted_dir=args.extracted,
processed_dir=args.processed,
sample_rate=args.sample_rate)
h36m.extract_tgz()
h36m.generate_cameras_file()
h36m.generate_annotations()
