def main(argv, common_opts):
args = parse_args(argv)
seed_all(12345)
init_algorithms(deterministic=True)
torch.set_grad_enabled(False)
device = common_opts['device']
if args.model:
# model = load_model(args.model).to(device).eval()
model = load_model(args.model).eval()
data_specs = model.data_specs
else:
model = None
data_specs = DataSpecs(
ImageSpecs(224,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(CanonicalSkeletonDesc, n_dims=3),
)
dataset = get_dataset(args.dataset, data_specs, use_aug=False)
app = MainGUIApp(dataset, device, model)
app.mainloop()
def __init__(self, data_dir, data_specs=None, use_aug=False, disable_mask_aug=False, \
without_image=True, human_height=2000, focal_diff=0, use_pcl=True, calculate_scale_from_2d=True, use_slant_compensation=True):
if data_specs is None:
data_specs = DataSpecs(
ImageSpecs(128,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(MpiInf3dhpSkeletonDesc, n_dims=3),
)
super().__init__(data_specs)
"""NEW"""
self.human_height = human_height
self.focal_diff = focal_diff
self.use_pcl = use_pcl
self.calculate_scale_from_2d = calculate_scale_from_2d
self.use_slant_compensation = use_slant_compensation
if not path.isdir(data_dir):
raise NotADirectoryError(data_dir)
metadata_files = sorted(
iglob(path.join(data_dir, 'S*', 'Seq*', 'metadata.h5')))
frame_refs = []
univ_scale_factors = {}
for metadata_file in metadata_files:
# match = re.match(r'.*S(\d+)/Seq(\d+)/metadata.h5', metadata_file)
match = re.match(r'.*S(\d+)\\Seq(\d+)\\metadata.h5', metadata_file)
subject_id = int(match.group(1))
sequence_id = int(match.group(2))
activity_ids = None
mat_annot_file = path.join(path.dirname(metadata_file),
'annot_data.mat')
if path.isfile(mat_annot_file):
with h5py.File(mat_annot_file, 'r') as f:
activity_ids = f['activity_annotation'][:].flatten(
).astype(int)
with h5py.File(metadata_file, 'r') as f:
keys = f['interesting_frames'].keys()
for key in keys:
camera_id = int(re.match(r'camera(\d)', key).group(1))
for frame_index in f['interesting_frames'][key]:
activity_id = None
if activity_ids is not None:
activity_id = activity_ids[frame_index]
frame_refs.append(
FrameRef(subject_id, sequence_id, camera_id,
frame_index, activity_id))
univ_scale_factors[(subject_id, sequence_id)] = f['scale'][0]
self.data_dir = data_dir
self.use_aug = use_aug
self.disable_mask_aug = disable_mask_aug
self.frame_refs = frame_refs
self.univ_scale_factors = univ_scale_factors
self.without_image = without_image
self.multicrop = False
def __init__(self, skel_desc, n_stages, axis_permutation, feature_extractor, pixelwise_loss):
super().__init__()
self.data_specs = DataSpecs(
ImageSpecs(256, mean=ImageSpecs.IMAGENET_MEAN, stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(skel_desc, n_dims=3),
)
self.pixelwise_loss = pixelwise_loss
self.inner = MargiPoseModelInner(skel_desc.n_joints, n_stages, axis_permutation,
feature_extractor)
self.xy_heatmaps = self.zy_heatmaps = self.xz_heatmaps = None
def test_mpi3d_val_subset(mpi3d_data_dir):
data_specs = DataSpecs(
ImageSpecs(256,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(CanonicalSkeletonDesc, n_dims=3),
)
dataset = MpiInf3dDataset(os.path.join(mpi3d_data_dir, 'val'), data_specs)
assert len(dataset) == 18561
example = dataset[0]
assert example['input'].shape == (3, 256, 256)
def __init__(self,
data_dir,
data_specs=None,
use_aug=False,
disable_mask_aug=False):
if data_specs is None:
data_specs = DataSpecs(
ImageSpecs(224,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(MpiInf3dhpSkeletonDesc, n_dims=3),
)
super().__init__(data_specs)
if not path.isdir(data_dir):
raise NotADirectoryError(data_dir)
metadata_files = sorted(
iglob(path.join(data_dir, 'S*', 'Seq*', 'metadata.h5')))
frame_refs = []
for metadata_file in metadata_files:
match = re.match(r'.*S(\d+)/Seq(\d+)/metadata.h5', metadata_file)
subject_id = int(match.group(1))
sequence_id = int(match.group(2))
activity_ids = None
mat_annot_file = path.join(path.dirname(metadata_file),
'annot_data.mat')
if path.isfile(mat_annot_file):
with h5py.File(mat_annot_file, 'r') as f:
activity_ids = f['activity_annotation'].value.flatten(
).astype(int)
with h5py.File(metadata_file, 'r') as f:
keys = f['interesting_frames'].keys()
for key in keys:
camera_id = int(re.match(r'camera(\d)', key).group(1))
for frame_index in f['interesting_frames'][key]:
activity_id = None
if activity_ids is not None:
activity_id = activity_ids[frame_index]
frame_refs.append(
FrameRef(subject_id, sequence_id, camera_id,
frame_index, activity_id))
self.data_dir = data_dir
self.use_aug = use_aug
self.disable_mask_aug = disable_mask_aug
self.frame_refs = frame_refs
self.without_image = False
self.multicrop = False
def __init__(self, skel_desc, pixelwise_loss):
super().__init__()
self.data_specs = DataSpecs(
ImageSpecs(256,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(skel_desc, n_dims=3),
)
self.pixelwise_loss = pixelwise_loss
resnet = resnet34(pretrained=True)
self.in_cnn = ResNetFeatureExtractor(resnet)
self.xy_hm_cnn = _XYCnn(resnet, skel_desc.n_joints)
self.zy_hm_cnn = _ChatterboxCnn(skel_desc.n_joints, shrink_width=True)
self.xz_hm_cnn = _ChatterboxCnn(skel_desc.n_joints, shrink_width=False)
def test_mpi3d_example_data(mpi3d_data_dir):
data_specs = DataSpecs(
ImageSpecs(256,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(CanonicalSkeletonDesc, n_dims=3),
)
dataset = MpiInf3dDataset(os.path.join(mpi3d_data_dir, 'val'), data_specs)
assert MpiInf3dDataset.preserve_root_joint_at_univ_scale == False
example = dataset[0]
image = example['input']
assert float(image.min()) == pytest.approx(-2.117904, rel=0, abs=1e-2)
assert float(image.max()) == pytest.approx(2.428571, rel=0, abs=1e-2)
joints = example['target'][..., :3]
assert_allclose(joints[0], [-0.025768, -0.649297, -0.039933],
rtol=0,
atol=1e-4)
def __init__(self,
data_dir,
data_specs=None,
subset='train',
use_aug=False,
max_length=None):
if data_specs is None:
data_specs = DataSpecs(
ImageSpecs(224,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(MpiiSkeletonDesc, n_dims=2),
)
super().__init__(data_specs)
self.subset = subset
self.use_aug = use_aug
self.mpii_data = MpiiData(data_dir)
self.example_ids = self.mpii_data.subset_indices(
self.subset)[:max_length]
def __init__(self, data_dir, human_height=2000, data_specs=None, subset='train', use_aug=False, max_length=None,
universal=False, focal_diff=0, without_image=False, use_pcl=True, calculate_scale_from_2d=False, use_slant_compensation=False,\
use_predicted_2D=False, predicted_input_dir=None, img_small_size=128, img_big_size=256):
if data_specs is None:
data_specs = DataSpecs(
ImageSpecs(img_small_size, mean=ImageSpecs.IMAGENET_MEAN, stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(H36MSkeletonDesc, n_dims=2),
)
super().__init__(data_specs)
if not path.isdir(data_dir):
raise NotADirectoryError(data_dir)
"""NEW"""
self.human_height = human_height
self.focal_diff = focal_diff
self.use_pcl = use_pcl
self.calculate_scale_from_2d = calculate_scale_from_2d
self.use_slant_compensation = use_slant_compensation
self.img_big_size = img_big_size
""""""
self.subset = subset
self.use_aug = use_aug
self.data_dir = data_dir
annot_files = sorted(iglob(path.join(data_dir, 'S*', '*', 'annot.h5')))
keys = ['pose/2d', 'pose/3d', 'pose/3d-univ', 'camera', 'frame',
'subject', 'action', 'subaction']
datasets = {}
self.camera_intrinsics = []
intrinsics_ds = 'intrinsics-univ' if universal else 'intrinsics'
self.predicted_input_dir = predicted_input_dir
self.use_predicted_2D = use_predicted_2D
if self.use_predicted_2D:
json_file_dict = {}
if self.predicted_input_dir is None:
predicted_pose_files = sorted(iglob(path.join(data_dir, 'hrnet_predictions', 'S*', '*', '*.json')))
else:
predicted_pose_files = sorted(iglob(path.join(predicted_input_dir, 'hrnet_predictions', 'S*', '*', '*.json')))
for json_file in predicted_pose_files:
with open(json_file) as f:
data = json.load(f)
for img in data:
frame_name = img['image_id']
directory_name = path.join(json_file, frame_name)
keypoints = img['keypoints']
json_file_dict[directory_name] = np.array(img['keypoints'], dtype=np.float32)
self.json_dict = json_file_dict
for annot_file in annot_files:
with h5py.File(annot_file) as annot:
for k in keys:
if k in datasets:
datasets[k].append(annot[k].value)
else:
datasets[k] = [annot[k].value]
cams = {}
for camera_id in annot[intrinsics_ds].keys():
alpha_x, x_0, alpha_y, y_0 = list(annot[intrinsics_ds][camera_id])
cams[int(camera_id)] = CameraIntrinsics.from_ccd_params(alpha_x, alpha_y, x_0, y_0)
for camera_id in annot['camera']:
self.camera_intrinsics.append(cams[camera_id])
datasets = {k: np.concatenate(v) for k, v in datasets.items()}
self.frame_ids = datasets['frame']
self.subject_ids = datasets['subject']
self.action_ids = datasets['action']
self.subaction_ids = datasets['subaction']
self.camera_ids = datasets['camera']
self.joint_3d = datasets['pose/3d-univ'] if universal else datasets['pose/3d']
self.joint_2d = datasets['pose/2d']
# Protocol #2
train_subjects = {1, 5, 6, 7, 8}
test_subjects = {9, 11}
train_ids = []
test_ids = []
for index, subject_id in enumerate(self.subject_ids):
if subject_id in train_subjects:
train_ids.append(index)
if subject_id in test_subjects:
test_ids.append(index)
if subset == 'trainval':
self.example_ids = np.array(train_ids, np.uint32)
elif subset == 'test':
self.example_ids = np.array(test_ids, np.uint32)
else:
raise Exception('Only trainval and test subsets are supported')
if max_length is not None:
self.example_ids = self.example_ids[:max_length]
self.without_image = without_image
self.multicrop = False
def __init__(self,
data_dir,
data_specs=None,
subset='train',
use_aug=False,
max_length=None,
universal=False):
if data_specs is None:
data_specs = DataSpecs(
ImageSpecs(224,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(H36MSkeletonDesc, n_dims=2),
)
super().__init__(data_specs)
if not path.isdir(data_dir):
raise NotADirectoryError(data_dir)
self.subset = subset
self.use_aug = use_aug
self.data_dir = data_dir
annot_files = sorted(iglob(path.join(data_dir, 'S*', '*', 'annot.h5')))
keys = [
'pose/2d', 'pose/3d', 'pose/3d-univ', 'camera', 'frame', 'subject',
'action', 'subaction'
]
datasets = {}
self.camera_intrinsics = []
intrinsics_ds = 'intrinsics-univ' if universal else 'intrinsics'
for annot_file in annot_files:
with h5py.File(annot_file) as annot:
for k in keys:
if k in datasets:
datasets[k].append(annot[k].value)
else:
datasets[k] = [annot[k].value]
cams = {}
for camera_id in annot[intrinsics_ds].keys():
alpha_x, x_0, alpha_y, y_0 = list(
annot[intrinsics_ds][camera_id])
cams[int(camera_id)] = CameraIntrinsics.from_ccd_params(
alpha_x, alpha_y, x_0, y_0)
for camera_id in annot['camera']:
self.camera_intrinsics.append(cams[camera_id])
datasets = {k: np.concatenate(v) for k, v in datasets.items()}
self.frame_ids = datasets['frame']
self.subject_ids = datasets['subject']
self.action_ids = datasets['action']
self.subaction_ids = datasets['subaction']
self.camera_ids = datasets['camera']
self.joint_3d = datasets['pose/3d-univ'] if universal else datasets[
'pose/3d']
self.joint_2d = datasets['pose/2d']
# Protocol #2
train_subjects = {1, 5, 6, 7, 8}
test_subjects = {9, 11}
train_ids = []
test_ids = []
for index, subject_id in enumerate(self.subject_ids):
if subject_id in train_subjects:
train_ids.append(index)
if subject_id in test_subjects:
test_ids.append(index)
if subset == 'trainval':
self.example_ids = np.array(train_ids, np.uint32)
elif subset == 'test':
self.example_ids = np.array(test_ids, np.uint32)
else:
raise Exception('Only trainval and test subsets are supported')
if max_length is not None:
self.example_ids = self.example_ids[:max_length]
self.without_image = False
self.multicrop = False