Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,165 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import os.path as osp
|
||||
from functools import wraps
|
||||
|
||||
import mmcv
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from mmpose.core import SimpleCamera
|
||||
|
||||
|
||||
def _keypoint_camera_to_world(keypoints,
|
||||
camera_params,
|
||||
image_name=None,
|
||||
dataset='Body3DH36MDataset'):
|
||||
"""Project 3D keypoints from the camera space to the world space.
|
||||
|
||||
Args:
|
||||
keypoints (np.ndarray): 3D keypoints in shape [..., 3]
|
||||
camera_params (dict): Parameters for all cameras.
|
||||
image_name (str): The image name to specify the camera.
|
||||
dataset (str): The dataset type, e.g., Body3DH36MDataset.
|
||||
"""
|
||||
cam_key = None
|
||||
if dataset == 'Body3DH36MDataset':
|
||||
subj, rest = osp.basename(image_name).split('_', 1)
|
||||
_, rest = rest.split('.', 1)
|
||||
camera, rest = rest.split('_', 1)
|
||||
cam_key = (subj, camera)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
camera = SimpleCamera(camera_params[cam_key])
|
||||
keypoints_world = keypoints.copy()
|
||||
keypoints_world[..., :3] = camera.camera_to_world(keypoints[..., :3])
|
||||
|
||||
return keypoints_world
|
||||
|
||||
|
||||
def _get_bbox_xywh(center, scale, w=200, h=200):
|
||||
w = w * scale
|
||||
h = h * scale
|
||||
x = center[0] - w / 2
|
||||
y = center[1] - h / 2
|
||||
return [x, y, w, h]
|
||||
|
||||
|
||||
def mmcv_track_func(func):
|
||||
|
||||
@wraps(func)
|
||||
def wrapped_func(args):
|
||||
return func(*args)
|
||||
|
||||
return wrapped_func
|
||||
|
||||
|
||||
@mmcv_track_func
|
||||
def _get_img_info(img_idx, img_name, img_root):
|
||||
try:
|
||||
im = Image.open(osp.join(img_root, img_name))
|
||||
w, h = im.size
|
||||
except: # noqa: E722
|
||||
return None
|
||||
|
||||
img = {
|
||||
'file_name': img_name,
|
||||
'height': h,
|
||||
'width': w,
|
||||
'id': img_idx + 1,
|
||||
}
|
||||
return img
|
||||
|
||||
|
||||
@mmcv_track_func
|
||||
def _get_ann(idx, kpt_2d, kpt_3d, center, scale, imgname, camera_params):
|
||||
bbox = _get_bbox_xywh(center, scale)
|
||||
kpt_3d = _keypoint_camera_to_world(kpt_3d, camera_params, imgname)
|
||||
|
||||
ann = {
|
||||
'id': idx + 1,
|
||||
'category_id': 1,
|
||||
'image_id': idx + 1,
|
||||
'iscrowd': 0,
|
||||
'bbox': bbox,
|
||||
'area': bbox[2] * bbox[3],
|
||||
'num_keypoints': 17,
|
||||
'keypoints': kpt_2d.reshape(-1).tolist(),
|
||||
'keypoints_3d': kpt_3d.reshape(-1).tolist()
|
||||
}
|
||||
|
||||
return ann
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--ann-file', type=str, default='tests/data/h36m/test_h36m_body3d.npz')
|
||||
parser.add_argument(
|
||||
'--camera-param-file', type=str, default='tests/data/h36m/cameras.pkl')
|
||||
parser.add_argument('--img-root', type=str, default='tests/data/h36m')
|
||||
parser.add_argument(
|
||||
'--out-file', type=str, default='tests/data/h36m/h36m_coco.json')
|
||||
parser.add_argument('--full-img-name', action='store_true')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
h36m_data = np.load(args.ann_file)
|
||||
h36m_camera_params = mmcv.load(args.camera_param_file)
|
||||
h36m_coco = {}
|
||||
|
||||
# categories
|
||||
h36m_cats = [{
|
||||
'supercategory':
|
||||
'person',
|
||||
'id':
|
||||
1,
|
||||
'name':
|
||||
'person',
|
||||
'keypoints': [
|
||||
'root (pelvis)', 'left_hip', 'left_knee', 'left_foot', 'right_hip',
|
||||
'right_knee', 'right_foot', 'spine', 'thorax', 'neck_base', 'head',
|
||||
'left_shoulder', 'left_elbow', 'left_wrist', 'right_shoulder',
|
||||
'right_elbow', 'right_wrist'
|
||||
],
|
||||
'skeleton': [[0, 1], [1, 2], [2, 3], [0, 4], [4, 5], [5, 6], [0, 7],
|
||||
[7, 8], [8, 9], [9, 10], [8, 11], [11, 12], [12, 13],
|
||||
[8, 14], [14, 15], [15, 16]],
|
||||
}]
|
||||
|
||||
# images
|
||||
imgnames = h36m_data['imgname']
|
||||
if not args.full_img_name:
|
||||
imgnames = [osp.basename(fn) for fn in imgnames]
|
||||
tasks = [(idx, fn, args.img_root) for idx, fn in enumerate(imgnames)]
|
||||
|
||||
h36m_imgs = mmcv.track_parallel_progress(_get_img_info, tasks, nproc=12)
|
||||
|
||||
# annotations
|
||||
kpts_2d = h36m_data['part']
|
||||
kpts_3d = h36m_data['S']
|
||||
centers = h36m_data['center']
|
||||
scales = h36m_data['scale']
|
||||
tasks = [(idx, ) + args + (h36m_camera_params, )
|
||||
for idx, args in enumerate(
|
||||
zip(kpts_2d, kpts_3d, centers, scales, imgnames))]
|
||||
|
||||
h36m_anns = mmcv.track_parallel_progress(_get_ann, tasks, nproc=12)
|
||||
|
||||
# remove invalid data
|
||||
h36m_imgs = [img for img in h36m_imgs if img is not None]
|
||||
h36m_img_ids = set([img['id'] for img in h36m_imgs])
|
||||
h36m_anns = [ann for ann in h36m_anns if ann['image_id'] in h36m_img_ids]
|
||||
|
||||
h36m_coco = {
|
||||
'categories': h36m_cats,
|
||||
'images': h36m_imgs,
|
||||
'annotations': h36m_anns,
|
||||
}
|
||||
|
||||
mmcv.dump(h36m_coco, args.out_file)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
|
||||
from scipy.io import loadmat
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Converting the predicted .mat file to .json file.')
|
||||
parser.add_argument('pred_mat_file', help='input prediction mat file.')
|
||||
parser.add_argument(
|
||||
'gt_json_file',
|
||||
help='input ground-truth json file to get the image name. '
|
||||
'Default: "data/mpii/mpii_val.json" ')
|
||||
parser.add_argument('output_json_file', help='output converted json file.')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def save_json(list_file, path):
|
||||
with open(path, 'w') as f:
|
||||
json.dump(list_file, f, indent=4)
|
||||
return 0
|
||||
|
||||
|
||||
def convert_mat(pred_mat_file, gt_json_file, output_json_file):
|
||||
res = loadmat(pred_mat_file)
|
||||
preds = res['preds']
|
||||
N = preds.shape[0]
|
||||
|
||||
with open(gt_json_file) as anno_file:
|
||||
anno = json.load(anno_file)
|
||||
|
||||
assert len(anno) == N
|
||||
|
||||
instance = {}
|
||||
|
||||
for pred, ann in zip(preds, anno):
|
||||
ann.pop('joints_vis')
|
||||
ann['joints'] = pred.tolist()
|
||||
|
||||
instance['annotations'] = anno
|
||||
instance['info'] = {}
|
||||
instance['info']['description'] = 'Converted MPII prediction.'
|
||||
instance['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
instance['info']['date_created'] = time.strftime('%Y/%m/%d',
|
||||
time.localtime())
|
||||
|
||||
save_json(instance, output_json_file)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
convert_mat(args.pred_mat_file, args.gt_json_file, args.output_json_file)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
+436
@@ -0,0 +1,436 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import warnings
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import xmltodict
|
||||
from xtcocotools.coco import COCO
|
||||
|
||||
np.random.seed(0)
|
||||
|
||||
|
||||
def list_all_files(root_dir, ext='.xml'):
|
||||
"""List all files in the root directory and all its sub directories.
|
||||
|
||||
:param root_dir: root directory
|
||||
:param ext: filename extension
|
||||
:return: list of files
|
||||
"""
|
||||
files = []
|
||||
file_list = os.listdir(root_dir)
|
||||
for i in range(0, len(file_list)):
|
||||
path = os.path.join(root_dir, file_list[i])
|
||||
if os.path.isdir(path):
|
||||
files.extend(list_all_files(path))
|
||||
if os.path.isfile(path):
|
||||
if path.lower().endswith(ext):
|
||||
files.append(path)
|
||||
return files
|
||||
|
||||
|
||||
def get_anno_info():
|
||||
keypoints_info = [
|
||||
'L_Eye',
|
||||
'R_Eye',
|
||||
'L_EarBase',
|
||||
'R_EarBase',
|
||||
'Nose',
|
||||
'Throat',
|
||||
'TailBase',
|
||||
'Withers',
|
||||
'L_F_Elbow',
|
||||
'R_F_Elbow',
|
||||
'L_B_Elbow',
|
||||
'R_B_Elbow',
|
||||
'L_F_Knee',
|
||||
'R_F_Knee',
|
||||
'L_B_Knee',
|
||||
'R_B_Knee',
|
||||
'L_F_Paw',
|
||||
'R_F_Paw',
|
||||
'L_B_Paw',
|
||||
'R_B_Paw',
|
||||
]
|
||||
skeleton_info = [[1, 2], [1, 3], [2, 4], [1, 5], [2, 5], [5, 6], [6, 8],
|
||||
[7, 8], [6, 9], [9, 13], [13, 17], [6, 10], [10, 14],
|
||||
[14, 18], [7, 11], [11, 15], [15, 19], [7, 12], [12, 16],
|
||||
[16, 20]]
|
||||
category_info = [{
|
||||
'supercategory': 'animal',
|
||||
'id': 1,
|
||||
'name': 'animal',
|
||||
'keypoints': keypoints_info,
|
||||
'skeleton': skeleton_info
|
||||
}]
|
||||
|
||||
return keypoints_info, skeleton_info, category_info
|
||||
|
||||
|
||||
def xml2coco_trainval(file_list, img_root, save_path, start_ann_id=0):
|
||||
"""Save annotations in coco-format.
|
||||
|
||||
:param file_list: list of data annotation files.
|
||||
:param img_root: the root dir to load images.
|
||||
:param save_path: the path to save transformed annotation file.
|
||||
:param start_ann_id: the starting point to count the annotation id.
|
||||
:param val_num: the number of annotated objects for validation.
|
||||
"""
|
||||
images = []
|
||||
annotations = []
|
||||
img_ids = []
|
||||
ann_ids = []
|
||||
|
||||
ann_id = start_ann_id
|
||||
|
||||
name2id = {
|
||||
'L_Eye': 0,
|
||||
'R_Eye': 1,
|
||||
'L_EarBase': 2,
|
||||
'R_EarBase': 3,
|
||||
'Nose': 4,
|
||||
'Throat': 5,
|
||||
'TailBase': 6,
|
||||
'Withers': 7,
|
||||
'L_F_Elbow': 8,
|
||||
'R_F_Elbow': 9,
|
||||
'L_B_Elbow': 10,
|
||||
'R_B_Elbow': 11,
|
||||
'L_F_Knee': 12,
|
||||
'R_F_Knee': 13,
|
||||
'L_B_Knee': 14,
|
||||
'R_B_Knee': 15,
|
||||
'L_F_Paw': 16,
|
||||
'R_F_Paw': 17,
|
||||
'L_B_Paw': 18,
|
||||
'R_B_Paw': 19
|
||||
}
|
||||
for file in file_list:
|
||||
data_anno = xmltodict.parse(open(file).read())['annotation']
|
||||
|
||||
img_id = int(data_anno['image'].split('_')[0] +
|
||||
data_anno['image'].split('_')[1])
|
||||
|
||||
if img_id not in img_ids:
|
||||
image_name = 'VOC2012/JPEGImages/' + data_anno['image'] + '.jpg'
|
||||
img = cv2.imread(os.path.join(img_root, image_name))
|
||||
|
||||
image = {}
|
||||
image['id'] = img_id
|
||||
image['file_name'] = image_name
|
||||
image['height'] = img.shape[0]
|
||||
image['width'] = img.shape[1]
|
||||
|
||||
images.append(image)
|
||||
img_ids.append(img_id)
|
||||
else:
|
||||
pass
|
||||
|
||||
keypoint_anno = data_anno['keypoints']['keypoint']
|
||||
assert len(keypoint_anno) == 20
|
||||
|
||||
keypoints = np.zeros([20, 3], dtype=np.float32)
|
||||
|
||||
for kpt_anno in keypoint_anno:
|
||||
keypoint_name = kpt_anno['@name']
|
||||
keypoint_id = name2id[keypoint_name]
|
||||
|
||||
visibility = int(kpt_anno['@visible'])
|
||||
|
||||
if visibility == 0:
|
||||
continue
|
||||
else:
|
||||
keypoints[keypoint_id, 0] = float(kpt_anno['@x'])
|
||||
keypoints[keypoint_id, 1] = float(kpt_anno['@y'])
|
||||
keypoints[keypoint_id, 2] = 2
|
||||
|
||||
anno = {}
|
||||
anno['keypoints'] = keypoints.reshape(-1).tolist()
|
||||
anno['image_id'] = img_id
|
||||
anno['id'] = ann_id
|
||||
anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))
|
||||
|
||||
visible_bounds = data_anno['visible_bounds']
|
||||
anno['bbox'] = [
|
||||
float(visible_bounds['@xmin']),
|
||||
float(visible_bounds['@ymin']),
|
||||
float(visible_bounds['@width']),
|
||||
float(visible_bounds['@height'])
|
||||
]
|
||||
anno['iscrowd'] = 0
|
||||
anno['area'] = float(anno['bbox'][2] * anno['bbox'][3])
|
||||
anno['category_id'] = 1
|
||||
|
||||
annotations.append(anno)
|
||||
ann_ids.append(ann_id)
|
||||
ann_id += 1
|
||||
|
||||
cocotype = {}
|
||||
|
||||
cocotype['info'] = {}
|
||||
cocotype['info'][
|
||||
'description'] = 'AnimalPose dataset Generated by MMPose Team'
|
||||
cocotype['info']['version'] = '1.0'
|
||||
cocotype['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
|
||||
time.localtime())
|
||||
|
||||
cocotype['images'] = images
|
||||
cocotype['annotations'] = annotations
|
||||
|
||||
keypoints_info, skeleton_info, category_info = get_anno_info()
|
||||
|
||||
cocotype['categories'] = category_info
|
||||
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
json.dump(cocotype, open(save_path, 'w'), indent=4)
|
||||
print('number of images:', len(img_ids))
|
||||
print('number of annotations:', len(ann_ids))
|
||||
print(f'done {save_path}')
|
||||
|
||||
|
||||
def xml2coco_test(file_list, img_root, save_path, start_ann_id=0):
|
||||
"""Save annotations in coco-format.
|
||||
|
||||
:param file_list: list of data annotation files.
|
||||
:param img_root: the root dir to load images.
|
||||
:param save_path: the path to save transformed annotation file.
|
||||
:param start_ann_id: the starting point to count the annotation id.
|
||||
"""
|
||||
images = []
|
||||
annotations = []
|
||||
img_ids = []
|
||||
ann_ids = []
|
||||
|
||||
ann_id = start_ann_id
|
||||
|
||||
name2id = {
|
||||
'L_eye': 0,
|
||||
'R_eye': 1,
|
||||
'L_ear': 2,
|
||||
'R_ear': 3,
|
||||
'Nose': 4,
|
||||
'Throat': 5,
|
||||
'Tail': 6,
|
||||
'withers': 7,
|
||||
'L_F_elbow': 8,
|
||||
'R_F_elbow': 9,
|
||||
'L_B_elbow': 10,
|
||||
'R_B_elbow': 11,
|
||||
'L_F_knee': 12,
|
||||
'R_F_knee': 13,
|
||||
'L_B_knee': 14,
|
||||
'R_B_knee': 15,
|
||||
'L_F_paw': 16,
|
||||
'R_F_paw': 17,
|
||||
'L_B_paw': 18,
|
||||
'R_B_paw': 19
|
||||
}
|
||||
|
||||
cat2id = {'cat': 1, 'cow': 2, 'dog': 3, 'horse': 4, 'sheep': 5}
|
||||
|
||||
for file in file_list:
|
||||
data_anno = xmltodict.parse(open(file).read())['annotation']
|
||||
|
||||
category_id = cat2id[data_anno['category']]
|
||||
|
||||
img_id = category_id * 1000 + int(
|
||||
re.findall(r'\d+', data_anno['image'])[0])
|
||||
|
||||
assert img_id not in img_ids
|
||||
|
||||
# prepare images
|
||||
image_name = os.path.join('animalpose_image_part2',
|
||||
data_anno['category'], data_anno['image'])
|
||||
img = cv2.imread(os.path.join(img_root, image_name))
|
||||
|
||||
image = {}
|
||||
image['id'] = img_id
|
||||
image['file_name'] = image_name
|
||||
image['height'] = img.shape[0]
|
||||
image['width'] = img.shape[1]
|
||||
|
||||
images.append(image)
|
||||
img_ids.append(img_id)
|
||||
|
||||
# prepare annotations
|
||||
keypoint_anno = data_anno['keypoints']['keypoint']
|
||||
keypoints = np.zeros([20, 3], dtype=np.float32)
|
||||
|
||||
for kpt_anno in keypoint_anno:
|
||||
keypoint_name = kpt_anno['@name']
|
||||
keypoint_id = name2id[keypoint_name]
|
||||
|
||||
visibility = int(kpt_anno['@visible'])
|
||||
|
||||
if visibility == 0:
|
||||
continue
|
||||
else:
|
||||
keypoints[keypoint_id, 0] = float(kpt_anno['@x'])
|
||||
keypoints[keypoint_id, 1] = float(kpt_anno['@y'])
|
||||
keypoints[keypoint_id, 2] = 2
|
||||
|
||||
anno = {}
|
||||
anno['keypoints'] = keypoints.reshape(-1).tolist()
|
||||
anno['image_id'] = img_id
|
||||
anno['id'] = ann_id
|
||||
anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))
|
||||
|
||||
visible_bounds = data_anno['visible_bounds']
|
||||
anno['bbox'] = [
|
||||
float(visible_bounds['@xmin']),
|
||||
float(visible_bounds['@xmax']
|
||||
), # typo in original xml: should be 'ymin'
|
||||
float(visible_bounds['@width']),
|
||||
float(visible_bounds['@height'])
|
||||
]
|
||||
anno['iscrowd'] = 0
|
||||
anno['area'] = float(anno['bbox'][2] * anno['bbox'][3])
|
||||
anno['category_id'] = 1
|
||||
|
||||
annotations.append(anno)
|
||||
ann_ids.append(ann_id)
|
||||
ann_id += 1
|
||||
|
||||
cocotype = {}
|
||||
|
||||
cocotype['info'] = {}
|
||||
cocotype['info'][
|
||||
'description'] = 'AnimalPose dataset Generated by MMPose Team'
|
||||
cocotype['info']['version'] = '1.0'
|
||||
cocotype['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
|
||||
time.localtime())
|
||||
|
||||
cocotype['images'] = images
|
||||
cocotype['annotations'] = annotations
|
||||
|
||||
keypoints_info, skeleton_info, category_info = get_anno_info()
|
||||
|
||||
cocotype['categories'] = category_info
|
||||
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
json.dump(cocotype, open(save_path, 'w'), indent=4)
|
||||
print('=========================================================')
|
||||
print('number of images:', len(img_ids))
|
||||
print('number of annotations:', len(ann_ids))
|
||||
print(f'done {save_path}')
|
||||
|
||||
|
||||
def split_train_val(work_dir, trainval_file, train_file, val_file,
|
||||
val_ann_num):
|
||||
"""Split train-val json file into training and validation files.
|
||||
|
||||
:param work_dir: path to load train-val json file, and save split files.
|
||||
:param trainval_file: The input json file combining both train and val.
|
||||
:param trainval_file: The output json file for training.
|
||||
:param trainval_file: The output json file for validation.
|
||||
:param val_ann_num: the number of validation annotations.
|
||||
"""
|
||||
|
||||
coco = COCO(os.path.join(work_dir, trainval_file))
|
||||
|
||||
img_list = list(coco.imgs.keys())
|
||||
np.random.shuffle(img_list)
|
||||
|
||||
count = 0
|
||||
|
||||
images_train = []
|
||||
images_val = []
|
||||
annotations_train = []
|
||||
annotations_val = []
|
||||
|
||||
for img_id in img_list:
|
||||
ann_ids = coco.getAnnIds(img_id)
|
||||
|
||||
if count + len(ann_ids) <= val_ann_num:
|
||||
# for validation
|
||||
count += len(ann_ids)
|
||||
images_val.append(coco.imgs[img_id])
|
||||
for ann_id in ann_ids:
|
||||
annotations_val.append(coco.anns[ann_id])
|
||||
|
||||
else:
|
||||
images_train.append(coco.imgs[img_id])
|
||||
for ann_id in ann_ids:
|
||||
annotations_train.append(coco.anns[ann_id])
|
||||
|
||||
if count == val_ann_num:
|
||||
print(f'We have found {count} annotations for validation.')
|
||||
else:
|
||||
warnings.warn(
|
||||
f'We only found {count} annotations, instead of {val_ann_num}.')
|
||||
|
||||
cocotype_train = {}
|
||||
cocotype_val = {}
|
||||
|
||||
keypoints_info, skeleton_info, category_info = get_anno_info()
|
||||
|
||||
cocotype_train['info'] = {}
|
||||
cocotype_train['info'][
|
||||
'description'] = 'AnimalPose dataset Generated by MMPose Team'
|
||||
cocotype_train['info']['version'] = '1.0'
|
||||
cocotype_train['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
cocotype_train['info']['date_created'] = time.strftime(
|
||||
'%Y/%m/%d', time.localtime())
|
||||
cocotype_train['images'] = images_train
|
||||
cocotype_train['annotations'] = annotations_train
|
||||
cocotype_train['categories'] = category_info
|
||||
|
||||
json.dump(
|
||||
cocotype_train,
|
||||
open(os.path.join(work_dir, train_file), 'w'),
|
||||
indent=4)
|
||||
print('=========================================================')
|
||||
print('number of images:', len(images_train))
|
||||
print('number of annotations:', len(annotations_train))
|
||||
print(f'done {train_file}')
|
||||
|
||||
cocotype_val['info'] = {}
|
||||
cocotype_val['info'][
|
||||
'description'] = 'AnimalPose dataset Generated by MMPose Team'
|
||||
cocotype_val['info']['version'] = '1.0'
|
||||
cocotype_val['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
cocotype_val['info']['date_created'] = time.strftime(
|
||||
'%Y/%m/%d', time.localtime())
|
||||
cocotype_val['images'] = images_val
|
||||
cocotype_val['annotations'] = annotations_val
|
||||
cocotype_val['categories'] = category_info
|
||||
|
||||
json.dump(
|
||||
cocotype_val, open(os.path.join(work_dir, val_file), 'w'), indent=4)
|
||||
print('=========================================================')
|
||||
print('number of images:', len(images_val))
|
||||
print('number of annotations:', len(annotations_val))
|
||||
print(f'done {val_file}')
|
||||
|
||||
|
||||
dataset_dir = 'data/animalpose/'
|
||||
|
||||
# We choose the images from PascalVOC for train + val
|
||||
# In total, train+val: 3608 images, 5117 annotations
|
||||
xml2coco_trainval(
|
||||
list_all_files(os.path.join(dataset_dir, 'PASCAL2011_animal_annotation')),
|
||||
dataset_dir,
|
||||
os.path.join(dataset_dir, 'annotations', 'animalpose_trainval.json'),
|
||||
start_ann_id=1000000)
|
||||
|
||||
# train: 2798 images, 4000 annotations
|
||||
# val: 810 images, 1117 annotations
|
||||
split_train_val(
|
||||
os.path.join(dataset_dir, 'annotations'),
|
||||
'animalpose_trainval.json',
|
||||
'animalpose_train.json',
|
||||
'animalpose_val.json',
|
||||
val_ann_num=1117)
|
||||
|
||||
# We choose the remaining 1000 images for test
|
||||
# 1000 images, 1000 annotations
|
||||
xml2coco_test(
|
||||
list_all_files(os.path.join(dataset_dir, 'animalpose_anno2')),
|
||||
dataset_dir,
|
||||
os.path.join(dataset_dir, 'annotations', 'animalpose_test.json'),
|
||||
start_ann_id=0)
|
||||
@@ -0,0 +1,97 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
mat_files = ['COFW_train_color.mat', 'COFW_test_color.mat']
|
||||
dataset_dir = 'data/cofw/'
|
||||
|
||||
image_root = os.path.join(dataset_dir, 'images/')
|
||||
annotation_root = os.path.join(dataset_dir, 'annotations/')
|
||||
|
||||
os.makedirs(image_root, exist_ok=True)
|
||||
os.makedirs(annotation_root, exist_ok=True)
|
||||
|
||||
cnt = 0
|
||||
for mat_file in mat_files:
|
||||
mat = h5py.File(os.path.join(dataset_dir, mat_file), 'r')
|
||||
|
||||
if 'train' in mat_file:
|
||||
imgs = mat['IsTr']
|
||||
pts = mat['phisTr']
|
||||
bboxes = mat['bboxesTr']
|
||||
is_train = True
|
||||
json_file = 'cofw_train.json'
|
||||
else:
|
||||
imgs = mat['IsT']
|
||||
pts = mat['phisT']
|
||||
bboxes = mat['bboxesT']
|
||||
is_train = False
|
||||
json_file = 'cofw_test.json'
|
||||
|
||||
images = []
|
||||
annotations = []
|
||||
|
||||
num = pts.shape[1]
|
||||
for idx in range(0, num):
|
||||
cnt += 1
|
||||
img = np.array(mat[imgs[0, idx]]).transpose()
|
||||
keypoints = pts[:, idx].reshape(3, -1).transpose()
|
||||
# 2 for valid and 1 for occlusion
|
||||
keypoints[:, 2] = 2 - keypoints[:, 2]
|
||||
# matlab 1-index to python 0-index
|
||||
keypoints[:, :2] -= 1
|
||||
bbox = bboxes[:, idx]
|
||||
|
||||
# check nonnegativity
|
||||
bbox[bbox < 0] = 0
|
||||
keypoints[keypoints < 0] = 0
|
||||
|
||||
image = {}
|
||||
image['id'] = cnt
|
||||
image['file_name'] = f'{str(cnt).zfill(6)}.jpg'
|
||||
image['height'] = img.shape[0]
|
||||
image['width'] = img.shape[1]
|
||||
cv2.imwrite(
|
||||
os.path.join(image_root, image['file_name']),
|
||||
cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
|
||||
images.append(image)
|
||||
|
||||
anno = {}
|
||||
anno['keypoints'] = keypoints.reshape(-1).tolist()
|
||||
anno['image_id'] = cnt
|
||||
anno['id'] = cnt
|
||||
anno['num_keypoints'] = len(keypoints) # all keypoints are labelled
|
||||
anno['bbox'] = bbox.tolist()
|
||||
anno['iscrowd'] = 0
|
||||
anno['area'] = anno['bbox'][2] * anno['bbox'][3]
|
||||
anno['category_id'] = 1
|
||||
|
||||
annotations.append(anno)
|
||||
|
||||
cocotype = {}
|
||||
|
||||
cocotype['info'] = {}
|
||||
cocotype['info']['description'] = 'COFW Generated by MMPose Team'
|
||||
cocotype['info']['version'] = '1.0'
|
||||
cocotype['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
|
||||
time.localtime())
|
||||
|
||||
cocotype['images'] = images
|
||||
cocotype['annotations'] = annotations
|
||||
cocotype['categories'] = [{
|
||||
'supercategory': 'person',
|
||||
'id': 1,
|
||||
'name': 'face',
|
||||
'keypoints': [],
|
||||
'skeleton': []
|
||||
}]
|
||||
|
||||
ann_path = os.path.join(annotation_root, json_file)
|
||||
json.dump(cocotype, open(ann_path, 'w'))
|
||||
print(f'done {ann_path}')
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
np.random.seed(0)
|
||||
|
||||
|
||||
def save_coco_anno(keypoints_all,
|
||||
annotated_all,
|
||||
imgs_all,
|
||||
keypoints_info,
|
||||
skeleton_info,
|
||||
dataset,
|
||||
img_root,
|
||||
save_path,
|
||||
start_img_id=0,
|
||||
start_ann_id=0):
|
||||
"""Save annotations in coco-format.
|
||||
|
||||
:param keypoints_all: keypoint annotations.
|
||||
:param annotated_all: images annotated or not.
|
||||
:param imgs_all: the array of images.
|
||||
:param keypoints_info: information about keypoint name.
|
||||
:param skeleton_info: information about skeleton connection.
|
||||
:param dataset: information about dataset name.
|
||||
:param img_root: the path to save images.
|
||||
:param save_path: the path to save transformed annotation file.
|
||||
:param start_img_id: the starting point to count the image id.
|
||||
:param start_ann_id: the starting point to count the annotation id.
|
||||
"""
|
||||
images = []
|
||||
annotations = []
|
||||
|
||||
img_id = start_img_id
|
||||
ann_id = start_ann_id
|
||||
|
||||
num_annotations, keypoints_num, _ = keypoints_all.shape
|
||||
|
||||
for i in range(num_annotations):
|
||||
img = imgs_all[i]
|
||||
keypoints = np.concatenate(
|
||||
[keypoints_all[i], annotated_all[i][:, None] * 2], axis=1)
|
||||
|
||||
min_x, min_y = np.min(keypoints[keypoints[:, 2] > 0, :2], axis=0)
|
||||
max_x, max_y = np.max(keypoints[keypoints[:, 2] > 0, :2], axis=0)
|
||||
|
||||
anno = {}
|
||||
anno['keypoints'] = keypoints.reshape(-1).tolist()
|
||||
anno['image_id'] = img_id
|
||||
anno['id'] = ann_id
|
||||
anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))
|
||||
anno['bbox'] = [
|
||||
float(min_x),
|
||||
float(min_y),
|
||||
float(max_x - min_x + 1),
|
||||
float(max_y - min_y + 1)
|
||||
]
|
||||
anno['iscrowd'] = 0
|
||||
anno['area'] = anno['bbox'][2] * anno['bbox'][3]
|
||||
anno['category_id'] = 1
|
||||
|
||||
annotations.append(anno)
|
||||
ann_id += 1
|
||||
|
||||
image = {}
|
||||
image['id'] = img_id
|
||||
image['file_name'] = f'{img_id}.jpg'
|
||||
image['height'] = img.shape[0]
|
||||
image['width'] = img.shape[1]
|
||||
|
||||
images.append(image)
|
||||
img_id += 1
|
||||
|
||||
cv2.imwrite(os.path.join(img_root, image['file_name']), img)
|
||||
|
||||
skeleton = np.concatenate(
|
||||
[np.arange(keypoints_num)[:, None], skeleton_info[:, 0][:, None]],
|
||||
axis=1) + 1
|
||||
skeleton = skeleton[skeleton.min(axis=1) > 0]
|
||||
|
||||
cocotype = {}
|
||||
|
||||
cocotype['info'] = {}
|
||||
cocotype['info'][
|
||||
'description'] = 'DeepPoseKit-Data Generated by MMPose Team'
|
||||
cocotype['info']['version'] = '1.0'
|
||||
cocotype['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
|
||||
time.localtime())
|
||||
|
||||
cocotype['images'] = images
|
||||
cocotype['annotations'] = annotations
|
||||
cocotype['categories'] = [{
|
||||
'supercategory': 'animal',
|
||||
'id': 1,
|
||||
'name': dataset,
|
||||
'keypoints': keypoints_info,
|
||||
'skeleton': skeleton.tolist()
|
||||
}]
|
||||
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
json.dump(cocotype, open(save_path, 'w'), indent=4)
|
||||
print('number of images:', img_id)
|
||||
print('number of annotations:', ann_id)
|
||||
print(f'done {save_path}')
|
||||
|
||||
|
||||
for dataset in ['fly', 'locust', 'zebra']:
|
||||
keypoints_info = []
|
||||
if dataset == 'fly':
|
||||
keypoints_info = [
|
||||
'head', 'eyeL', 'eyeR', 'neck', 'thorax', 'abdomen', 'forelegR1',
|
||||
'forelegR2', 'forelegR3', 'forelegR4', 'midlegR1', 'midlegR2',
|
||||
'midlegR3', 'midlegR4', 'hindlegR1', 'hindlegR2', 'hindlegR3',
|
||||
'hindlegR4', 'forelegL1', 'forelegL2', 'forelegL3', 'forelegL4',
|
||||
'midlegL1', 'midlegL2', 'midlegL3', 'midlegL4', 'hindlegL1',
|
||||
'hindlegL2', 'hindlegL3', 'hindlegL4', 'wingL', 'wingR'
|
||||
]
|
||||
elif dataset == 'locust':
|
||||
keypoints_info = [
|
||||
'head', 'neck', 'thorax', 'abdomen1', 'abdomen2', 'anttipL',
|
||||
'antbaseL', 'eyeL', 'forelegL1', 'forelegL2', 'forelegL3',
|
||||
'forelegL4', 'midlegL1', 'midlegL2', 'midlegL3', 'midlegL4',
|
||||
'hindlegL1', 'hindlegL2', 'hindlegL3', 'hindlegL4', 'anttipR',
|
||||
'antbaseR', 'eyeR', 'forelegR1', 'forelegR2', 'forelegR3',
|
||||
'forelegR4', 'midlegR1', 'midlegR2', 'midlegR3', 'midlegR4',
|
||||
'hindlegR1', 'hindlegR2', 'hindlegR3', 'hindlegR4'
|
||||
]
|
||||
elif dataset == 'zebra':
|
||||
keypoints_info = [
|
||||
'snout', 'head', 'neck', 'forelegL1', 'forelegR1', 'hindlegL1',
|
||||
'hindlegR1', 'tailbase', 'tailtip'
|
||||
]
|
||||
else:
|
||||
NotImplementedError()
|
||||
|
||||
dataset_dir = f'data/DeepPoseKit-Data/datasets/{dataset}'
|
||||
|
||||
with h5py.File(
|
||||
os.path.join(dataset_dir, 'annotation_data_release.h5'), 'r') as f:
|
||||
# List all groups
|
||||
annotations = np.array(f['annotations'])
|
||||
annotated = np.array(f['annotated'])
|
||||
images = np.array(f['images'])
|
||||
skeleton_info = np.array(f['skeleton'])
|
||||
|
||||
annotation_num, kpt_num, _ = annotations.shape
|
||||
|
||||
data_list = np.arange(0, annotation_num)
|
||||
np.random.shuffle(data_list)
|
||||
|
||||
val_data_num = annotation_num // 10
|
||||
train_data_num = annotation_num - val_data_num
|
||||
|
||||
train_list = data_list[0:train_data_num]
|
||||
val_list = data_list[train_data_num:]
|
||||
|
||||
img_root = os.path.join(dataset_dir, 'images')
|
||||
os.makedirs(img_root, exist_ok=True)
|
||||
|
||||
save_coco_anno(
|
||||
annotations[train_list], annotated[train_list], images[train_list],
|
||||
keypoints_info, skeleton_info, dataset, img_root,
|
||||
os.path.join(dataset_dir, 'annotations', f'{dataset}_train.json'))
|
||||
save_coco_anno(
|
||||
annotations[val_list],
|
||||
annotated[val_list],
|
||||
images[val_list],
|
||||
keypoints_info,
|
||||
skeleton_info,
|
||||
dataset,
|
||||
img_root,
|
||||
os.path.join(dataset_dir, 'annotations', f'{dataset}_test.json'),
|
||||
start_img_id=train_data_num,
|
||||
start_ann_id=train_data_num)
|
||||
+182
@@ -0,0 +1,182 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
np.random.seed(0)
|
||||
|
||||
|
||||
def get_poly_area(x, y):
|
||||
"""Calculate area of polygon given (x,y) coordinates (Shoelace formula)
|
||||
|
||||
:param x: np.ndarray(N, )
|
||||
:param y: np.ndarray(N, )
|
||||
:return: area
|
||||
"""
|
||||
return float(0.5 *
|
||||
np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1))))
|
||||
|
||||
|
||||
def get_seg_area(segmentations):
|
||||
area = 0
|
||||
for segmentation in segmentations:
|
||||
area += get_poly_area(segmentation[:, 0], segmentation[:, 1])
|
||||
return area
|
||||
|
||||
|
||||
def save_coco_anno(data_annotation,
|
||||
img_root,
|
||||
save_path,
|
||||
start_img_id=0,
|
||||
start_ann_id=0,
|
||||
kpt_num=17):
|
||||
"""Save annotations in coco-format.
|
||||
|
||||
:param data_annotation: list of data annotation.
|
||||
:param img_root: the root dir to load images.
|
||||
:param save_path: the path to save transformed annotation file.
|
||||
:param start_img_id: the starting point to count the image id.
|
||||
:param start_ann_id: the starting point to count the annotation id.
|
||||
:param kpt_num: the number of keypoint.
|
||||
"""
|
||||
images = []
|
||||
annotations = []
|
||||
|
||||
img_id = start_img_id
|
||||
ann_id = start_ann_id
|
||||
|
||||
for i in range(0, len(data_annotation)):
|
||||
data_anno = data_annotation[i]
|
||||
image_name = data_anno[0]
|
||||
|
||||
img = cv2.imread(os.path.join(img_root, image_name))
|
||||
|
||||
kp_string = data_anno[1]
|
||||
kps = json.loads(kp_string)
|
||||
|
||||
seg_string = data_anno[2]
|
||||
segs = json.loads(seg_string)
|
||||
|
||||
for kp, seg in zip(kps, segs):
|
||||
keypoints = np.zeros([kpt_num, 3])
|
||||
for ind, p in enumerate(kp):
|
||||
if p['position'] is None:
|
||||
continue
|
||||
else:
|
||||
keypoints[ind, 0] = p['position'][0]
|
||||
keypoints[ind, 1] = p['position'][1]
|
||||
keypoints[ind, 2] = 2
|
||||
|
||||
segmentations = []
|
||||
|
||||
max_x = -1
|
||||
max_y = -1
|
||||
min_x = 999999
|
||||
min_y = 999999
|
||||
for segm in seg:
|
||||
if len(segm['segment']) == 0:
|
||||
continue
|
||||
|
||||
segmentation = np.array(segm['segment'])
|
||||
segmentations.append(segmentation)
|
||||
|
||||
_max_x, _max_y = segmentation.max(0)
|
||||
_min_x, _min_y = segmentation.min(0)
|
||||
|
||||
max_x = max(max_x, _max_x)
|
||||
max_y = max(max_y, _max_y)
|
||||
min_x = min(min_x, _min_x)
|
||||
min_y = min(min_y, _min_y)
|
||||
|
||||
anno = {}
|
||||
anno['keypoints'] = keypoints.reshape(-1).tolist()
|
||||
anno['image_id'] = img_id
|
||||
anno['id'] = ann_id
|
||||
anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))
|
||||
anno['bbox'] = [
|
||||
float(min_x),
|
||||
float(min_y),
|
||||
float(max_x - min_x + 1),
|
||||
float(max_y - min_y + 1)
|
||||
]
|
||||
anno['iscrowd'] = 0
|
||||
anno['area'] = get_seg_area(segmentations)
|
||||
anno['category_id'] = 1
|
||||
anno['segmentation'] = [
|
||||
seg.reshape(-1).tolist() for seg in segmentations
|
||||
]
|
||||
|
||||
annotations.append(anno)
|
||||
ann_id += 1
|
||||
|
||||
image = {}
|
||||
image['id'] = img_id
|
||||
image['file_name'] = image_name
|
||||
image['height'] = img.shape[0]
|
||||
image['width'] = img.shape[1]
|
||||
|
||||
images.append(image)
|
||||
img_id += 1
|
||||
|
||||
cocotype = {}
|
||||
|
||||
cocotype['info'] = {}
|
||||
cocotype['info']['description'] = 'MacaquePose Generated by MMPose Team'
|
||||
cocotype['info']['version'] = '1.0'
|
||||
cocotype['info']['year'] = time.strftime('%Y', time.localtime())
|
||||
cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
|
||||
time.localtime())
|
||||
|
||||
cocotype['images'] = images
|
||||
cocotype['annotations'] = annotations
|
||||
cocotype['categories'] = [{
|
||||
'supercategory':
|
||||
'animal',
|
||||
'id':
|
||||
1,
|
||||
'name':
|
||||
'macaque',
|
||||
'keypoints': [
|
||||
'nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
|
||||
'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',
|
||||
'left_wrist', 'right_wrist', 'left_hip', 'right_hip', 'left_knee',
|
||||
'right_knee', 'left_ankle', 'right_ankle'
|
||||
],
|
||||
'skeleton': [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12],
|
||||
[7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3],
|
||||
[1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]
|
||||
}]
|
||||
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
json.dump(cocotype, open(save_path, 'w'), indent=4)
|
||||
print('number of images:', img_id)
|
||||
print('number of annotations:', ann_id)
|
||||
print(f'done {save_path}')
|
||||
|
||||
|
||||
dataset_dir = '/data/macaque/'
|
||||
with open(os.path.join(dataset_dir, 'annotations.csv'), 'r') as fp:
|
||||
data_annotation_all = list(csv.reader(fp, delimiter=','))[1:]
|
||||
|
||||
np.random.shuffle(data_annotation_all)
|
||||
|
||||
data_annotation_train = data_annotation_all[0:12500]
|
||||
data_annotation_val = data_annotation_all[12500:]
|
||||
|
||||
img_root = os.path.join(dataset_dir, 'images')
|
||||
save_coco_anno(
|
||||
data_annotation_train,
|
||||
img_root,
|
||||
os.path.join(dataset_dir, 'annotations', 'macaque_train.json'),
|
||||
kpt_num=17)
|
||||
save_coco_anno(
|
||||
data_annotation_val,
|
||||
img_root,
|
||||
os.path.join(dataset_dir, 'annotations', 'macaque_test.json'),
|
||||
start_img_id=12500,
|
||||
start_ann_id=15672,
|
||||
kpt_num=17)
|
||||
@@ -0,0 +1,417 @@
|
||||
# -----------------------------------------------------------------------------
|
||||
# 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()
|
||||
+359
@@ -0,0 +1,359 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import os
|
||||
import pickle
|
||||
import shutil
|
||||
from os.path import join
|
||||
|
||||
import cv2
|
||||
import h5py
|
||||
import mmcv
|
||||
import numpy as np
|
||||
from scipy.io import loadmat
|
||||
|
||||
train_subjects = [i for i in range(1, 9)]
|
||||
test_subjects = [i for i in range(1, 7)]
|
||||
train_seqs = [1, 2]
|
||||
train_cams = [0, 1, 2, 4, 5, 6, 7, 8]
|
||||
train_frame_nums = {
|
||||
(1, 1): 6416,
|
||||
(1, 2): 12430,
|
||||
(2, 1): 6502,
|
||||
(2, 2): 6081,
|
||||
(3, 1): 12488,
|
||||
(3, 2): 12283,
|
||||
(4, 1): 6171,
|
||||
(4, 2): 6675,
|
||||
(5, 1): 12820,
|
||||
(5, 2): 12312,
|
||||
(6, 1): 6188,
|
||||
(6, 2): 6145,
|
||||
(7, 1): 6239,
|
||||
(7, 2): 6320,
|
||||
(8, 1): 6468,
|
||||
(8, 2): 6054
|
||||
}
|
||||
test_frame_nums = {1: 6151, 2: 6080, 3: 5838, 4: 6007, 5: 320, 6: 492}
|
||||
train_img_size = (2048, 2048)
|
||||
root_index = 14
|
||||
joints_17 = [7, 5, 14, 15, 16, 9, 10, 11, 23, 24, 25, 18, 19, 20, 4, 3, 6]
|
||||
|
||||
|
||||
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 get_annotations(joints_2d, joints_3d, scale_factor=1.2):
|
||||
"""Get annotations, including centers, scales, joints_2d and joints_3d.
|
||||
|
||||
Args:
|
||||
joints_2d: 2D joint coordinates in shape [N, K, 2], where N is the
|
||||
frame number, K is the joint number.
|
||||
joints_3d: 3D joint coordinates in shape [N, K, 3], where N is the
|
||||
frame number, K is the joint number.
|
||||
scale_factor: Scale factor of bounding box. Default: 1.2.
|
||||
Returns:
|
||||
centers (ndarray): [N, 2]
|
||||
scales (ndarray): [N,]
|
||||
joints_2d (ndarray): [N, K, 3]
|
||||
joints_3d (ndarray): [N, K, 4]
|
||||
"""
|
||||
# calculate joint visibility
|
||||
visibility = (joints_2d[:, :, 0] >= 0) * \
|
||||
(joints_2d[:, :, 0] < train_img_size[0]) * \
|
||||
(joints_2d[:, :, 1] >= 0) * \
|
||||
(joints_2d[:, :, 1] < train_img_size[1])
|
||||
visibility = np.array(visibility, dtype=np.float32)[:, :, None]
|
||||
joints_2d = np.concatenate([joints_2d, visibility], axis=-1)
|
||||
joints_3d = np.concatenate([joints_3d, visibility], axis=-1)
|
||||
|
||||
# calculate bounding boxes
|
||||
bboxes = np.stack([
|
||||
np.min(joints_2d[:, :, 0], axis=1),
|
||||
np.min(joints_2d[:, :, 1], axis=1),
|
||||
np.max(joints_2d[:, :, 0], axis=1),
|
||||
np.max(joints_2d[:, :, 1], axis=1)
|
||||
],
|
||||
axis=1)
|
||||
centers = np.stack([(bboxes[:, 0] + bboxes[:, 2]) / 2,
|
||||
(bboxes[:, 1] + bboxes[:, 3]) / 2],
|
||||
axis=1)
|
||||
scales = scale_factor * np.max(bboxes[:, 2:] - bboxes[:, :2], axis=1) / 200
|
||||
|
||||
return centers, scales, joints_2d, joints_3d
|
||||
|
||||
|
||||
def load_trainset(data_root, out_dir):
|
||||
"""Load training data, create annotation file and camera file.
|
||||
Args:
|
||||
data_root: Directory of dataset, which is organized in the following
|
||||
hierarchy:
|
||||
data_root
|
||||
|-- train
|
||||
|-- S1
|
||||
|-- Seq1
|
||||
|-- Seq2
|
||||
|-- S2
|
||||
|-- ...
|
||||
|-- test
|
||||
|-- TS1
|
||||
|-- TS2
|
||||
|-- ...
|
||||
out_dir: Directory to save annotation file.
|
||||
"""
|
||||
_imgnames = []
|
||||
_centers = []
|
||||
_scales = []
|
||||
_joints_2d = []
|
||||
_joints_3d = []
|
||||
cameras = {}
|
||||
|
||||
img_dir = join(out_dir, 'images')
|
||||
os.makedirs(img_dir, exist_ok=True)
|
||||
annot_dir = join(out_dir, 'annotations')
|
||||
os.makedirs(annot_dir, exist_ok=True)
|
||||
|
||||
for subj in train_subjects:
|
||||
for seq in train_seqs:
|
||||
seq_path = join(data_root, 'train', f'S{subj}', f'Seq{seq}')
|
||||
num_frames = train_frame_nums[(subj, seq)]
|
||||
|
||||
# load camera parametres
|
||||
camera_file = join(seq_path, 'camera.calibration')
|
||||
with open(camera_file, 'r') as fin:
|
||||
lines = fin.readlines()
|
||||
for cam in train_cams:
|
||||
K = [float(s) for s in lines[cam * 7 + 5][11:-2].split()]
|
||||
f = np.array([[K[0]], [K[5]]])
|
||||
c = np.array([[K[2]], [K[6]]])
|
||||
RT = np.array(
|
||||
[float(s) for s in lines[cam * 7 + 6][11:-2].split()])
|
||||
RT = np.reshape(RT, (4, 4))
|
||||
R = RT[:3, :3]
|
||||
# convert unit from millimeter to meter
|
||||
T = RT[:3, 3:] * 0.001
|
||||
size = [int(s) for s in lines[cam * 7 + 3][14:].split()]
|
||||
w, h = size
|
||||
cam_param = dict(
|
||||
R=R, T=T, c=c, f=f, w=w, h=h, name=f'train_cam_{cam}')
|
||||
cameras[f'S{subj}_Seq{seq}_Cam{cam}'] = cam_param
|
||||
|
||||
# load annotations
|
||||
annot_file = os.path.join(seq_path, 'annot.mat')
|
||||
annot2 = loadmat(annot_file)['annot2']
|
||||
annot3 = loadmat(annot_file)['annot3']
|
||||
for cam in train_cams:
|
||||
# load 2D and 3D annotations
|
||||
joints_2d = np.reshape(annot2[cam][0][:num_frames],
|
||||
(num_frames, 28, 2))[:, joints_17]
|
||||
joints_3d = np.reshape(annot3[cam][0][:num_frames],
|
||||
(num_frames, 28, 3))[:, joints_17]
|
||||
joints_3d = joints_3d * 0.001
|
||||
centers, scales, joints_2d, joints_3d = get_annotations(
|
||||
joints_2d, joints_3d)
|
||||
_centers.append(centers)
|
||||
_scales.append(scales)
|
||||
_joints_2d.append(joints_2d)
|
||||
_joints_3d.append(joints_3d)
|
||||
|
||||
# extract frames from video
|
||||
video_path = join(seq_path, 'imageSequence',
|
||||
f'video_{cam}.avi')
|
||||
video = mmcv.VideoReader(video_path)
|
||||
for i in mmcv.track_iter_progress(range(num_frames)):
|
||||
img = video.read()
|
||||
if img is None:
|
||||
break
|
||||
imgname = f'S{subj}_Seq{seq}_Cam{cam}_{i+1:06d}.jpg'
|
||||
_imgnames.append(imgname)
|
||||
cv2.imwrite(join(img_dir, imgname), img)
|
||||
|
||||
_imgnames = np.array(_imgnames)
|
||||
_centers = np.concatenate(_centers)
|
||||
_scales = np.concatenate(_scales)
|
||||
_joints_2d = np.concatenate(_joints_2d)
|
||||
_joints_3d = np.concatenate(_joints_3d)
|
||||
|
||||
out_file = join(annot_dir, 'mpi_inf_3dhp_train.npz')
|
||||
np.savez(
|
||||
out_file,
|
||||
imgname=_imgnames,
|
||||
center=_centers,
|
||||
scale=_scales,
|
||||
part=_joints_2d,
|
||||
S=_joints_3d)
|
||||
print(f'Create annotation file for trainset: {out_file}. '
|
||||
f'{len(_imgnames)} samples in total.')
|
||||
|
||||
out_file = join(annot_dir, 'cameras_train.pkl')
|
||||
with open(out_file, 'wb') as fout:
|
||||
pickle.dump(cameras, fout)
|
||||
print(f'Create camera file for trainset: {out_file}.')
|
||||
|
||||
# get `mean` and `std` of pose data
|
||||
_joints_3d = _joints_3d[..., :3] # remove visibility
|
||||
mean_3d, std_3d = get_pose_stats(_joints_3d)
|
||||
|
||||
_joints_2d = _joints_2d[..., :2] # remove visibility
|
||||
mean_2d, std_2d = get_pose_stats(_joints_2d)
|
||||
|
||||
# centered around root
|
||||
_joints_3d_rel = _joints_3d - _joints_3d[..., root_index:root_index + 1, :]
|
||||
mean_3d_rel, std_3d_rel = get_pose_stats(_joints_3d_rel)
|
||||
mean_3d_rel[root_index] = mean_3d[root_index]
|
||||
std_3d_rel[root_index] = std_3d[root_index]
|
||||
|
||||
_joints_2d_rel = _joints_2d - _joints_2d[..., root_index:root_index + 1, :]
|
||||
mean_2d_rel, std_2d_rel = get_pose_stats(_joints_2d_rel)
|
||||
mean_2d_rel[root_index] = mean_2d[root_index]
|
||||
std_2d_rel[root_index] = std_2d[root_index]
|
||||
|
||||
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(annot_dir, f'{name}.pkl')
|
||||
with open(out_file, 'wb') as f:
|
||||
pickle.dump(stat_dict, f)
|
||||
print(f'Create statistic data file: {out_file}')
|
||||
|
||||
|
||||
def load_testset(data_root, out_dir, valid_only=True):
|
||||
"""Load testing data, create annotation file and camera file.
|
||||
|
||||
Args:
|
||||
data_root: Directory of dataset.
|
||||
out_dir: Directory to save annotation file.
|
||||
valid_only: Only keep frames with valid_label == 1.
|
||||
"""
|
||||
_imgnames = []
|
||||
_centers = []
|
||||
_scales = []
|
||||
_joints_2d = []
|
||||
_joints_3d = []
|
||||
cameras = {}
|
||||
|
||||
img_dir = join(out_dir, 'images')
|
||||
os.makedirs(img_dir, exist_ok=True)
|
||||
annot_dir = join(out_dir, 'annotations')
|
||||
os.makedirs(annot_dir, exist_ok=True)
|
||||
|
||||
for subj in test_subjects:
|
||||
subj_path = join(data_root, 'test', f'TS{subj}')
|
||||
num_frames = test_frame_nums[subj]
|
||||
|
||||
# load annotations
|
||||
annot_file = os.path.join(subj_path, 'annot_data.mat')
|
||||
with h5py.File(annot_file, 'r') as fin:
|
||||
annot2 = np.array(fin['annot2']).reshape((-1, 17, 2))
|
||||
annot3 = np.array(fin['annot3']).reshape((-1, 17, 3))
|
||||
valid = np.array(fin['valid_frame']).reshape(-1)
|
||||
|
||||
# manually estimate camera intrinsics
|
||||
fx, cx = np.linalg.lstsq(
|
||||
annot3[:, :, [0, 2]].reshape((-1, 2)),
|
||||
(annot2[:, :, 0] * annot3[:, :, 2]).reshape(-1, 1),
|
||||
rcond=None)[0].flatten()
|
||||
fy, cy = np.linalg.lstsq(
|
||||
annot3[:, :, [1, 2]].reshape((-1, 2)),
|
||||
(annot2[:, :, 1] * annot3[:, :, 2]).reshape(-1, 1),
|
||||
rcond=None)[0].flatten()
|
||||
if subj <= 4:
|
||||
w, h = 2048, 2048
|
||||
else:
|
||||
w, h = 1920, 1080
|
||||
cameras[f'TS{subj}'] = dict(
|
||||
c=np.array([[cx], [cy]]),
|
||||
f=np.array([[fx], [fy]]),
|
||||
w=w,
|
||||
h=h,
|
||||
name=f'test_cam_{subj}')
|
||||
|
||||
# get annotations
|
||||
if valid_only:
|
||||
valid_frames = np.nonzero(valid)[0]
|
||||
else:
|
||||
valid_frames = np.arange(num_frames)
|
||||
joints_2d = annot2[valid_frames, :, :]
|
||||
joints_3d = annot3[valid_frames, :, :] * 0.001
|
||||
|
||||
centers, scales, joints_2d, joints_3d = get_annotations(
|
||||
joints_2d, joints_3d)
|
||||
_centers.append(centers)
|
||||
_scales.append(scales)
|
||||
_joints_2d.append(joints_2d)
|
||||
_joints_3d.append(joints_3d)
|
||||
|
||||
# copy and rename images
|
||||
for i in valid_frames:
|
||||
imgname = f'TS{subj}_{i+1:06d}.jpg'
|
||||
shutil.copyfile(
|
||||
join(subj_path, 'imageSequence', f'img_{i+1:06d}.jpg'),
|
||||
join(img_dir, imgname))
|
||||
_imgnames.append(imgname)
|
||||
|
||||
_imgnames = np.array(_imgnames)
|
||||
_centers = np.concatenate(_centers)
|
||||
_scales = np.concatenate(_scales)
|
||||
_joints_2d = np.concatenate(_joints_2d)
|
||||
_joints_3d = np.concatenate(_joints_3d)
|
||||
|
||||
if valid_only:
|
||||
out_file = join(annot_dir, 'mpi_inf_3dhp_test_valid.npz')
|
||||
else:
|
||||
out_file = join(annot_dir, 'mpi_inf_3dhp_test_all.npz')
|
||||
np.savez(
|
||||
out_file,
|
||||
imgname=_imgnames,
|
||||
center=_centers,
|
||||
scale=_scales,
|
||||
part=_joints_2d,
|
||||
S=_joints_3d)
|
||||
print(f'Create annotation file for testset: {out_file}. '
|
||||
f'{len(_imgnames)} samples in total.')
|
||||
|
||||
out_file = join(annot_dir, 'cameras_test.pkl')
|
||||
with open(out_file, 'wb') as fout:
|
||||
pickle.dump(cameras, fout)
|
||||
print(f'Create camera file for testset: {out_file}.')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('data_root', type=str, help='data root')
|
||||
parser.add_argument(
|
||||
'out_dir', type=str, help='directory to save annotation files.')
|
||||
args = parser.parse_args()
|
||||
data_root = args.data_root
|
||||
out_dir = args.out_dir
|
||||
|
||||
load_trainset(data_root, out_dir)
|
||||
load_testset(data_root, out_dir, valid_only=True)
|
||||
Reference in New Issue
Block a user