print('Loading 2D detections...')
keypoints = np.load('/home/kjakkala/tmp.npz', allow_pickle=True)
keypoints_metadata = keypoints['metadata'].item()
keypoints_symmetry = keypoints_metadata['keypoints_symmetry']
kps_left, kps_right = list(keypoints_symmetry[0]), list(keypoints_symmetry[1])
joints_left, joints_right = list(dataset.skeleton().joints_left()), list(
dataset.skeleton().joints_right())
print(joints_left, joints_right)
keypoints = keypoints['positions_2d'].item()
for subject in keypoints.keys():
for action in keypoints[subject]:
for cam_idx, kps in enumerate(keypoints[subject][action]):
# Normalize camera frame
cam = dataset.cameras()[subject][cam_idx]
kps[..., :2] = normalize_screen_coordinates(kps[..., :2],
w=cam['res_w'],
h=cam['res_h'])
print(kps.shape)
keypoints[subject][action][cam_idx] = kps
def fetch():
out_poses_2d = []
for subject in keypoints.keys():
for action in keypoints[subject].keys():
poses_2d = keypoints[subject][action]
for i in range(len(poses_2d)): # Iterate across cameras
out_poses_2d.append(poses_2d[i])
return out_poses_2d
for cam_idx in range(len(keypoints[subject][action])):
# We check for >= instead of == because some videos in H3.6M contain extra frames
mocap_length = dataset[subject][action]['positions_3d'][
cam_idx].shape[0]
assert keypoints[subject][action][cam_idx].shape[0] >= mocap_length
if keypoints[subject][action][cam_idx].shape[0] > mocap_length:
# Shorten sequence
keypoints[subject][action][cam_idx] = keypoints[subject][
action][cam_idx][:mocap_length]
assert len(keypoints[subject][action]) == len(
dataset[subject][action]['positions_3d'])
print(dataset.cameras())
#NORMALIZE THE KEYPOINTS FOR THE RES
for subject in keypoints.keys(
): #should just be one subject if one video (i.e. '../vids/output.mp4')
print("SUBJECT: ", subject)
for action in keypoints[
subject]: #should just be one action, which is 'custom'
print("ACTION: ", action)
for cam_idx, kps in enumerate(
keypoints[subject][action]
): #each kps should be the Fx17x2 array of coordinates, where F is num frames
print("INDEX ", cam_idx, "KPS: ", kps)
print(dataset.cameras()[subject][cam_idx])
prediction += prediction_traj
if args.viz_export is not None:
print('Exporting joint positions to', args.viz_export)
# Predictions are in camera space
np.save(args.viz_export, prediction)
if args.viz_output is not None:
if ground_truth is not None:
# Reapply trajectory
trajectory = ground_truth[:, :1]
ground_truth[:, 1:] += trajectory
prediction += trajectory
# Invert camera transformation
cam = dataset.cameras()[args.viz_subject][args.viz_camera]
if ground_truth is not None:
prediction = camera_to_world(
prediction, R=cam['orientation'], t=cam['translation'])
ground_truth = camera_to_world(
ground_truth, R=cam['orientation'], t=cam['translation'])
else:
# If the ground truth is not available, take the camera extrinsic params from a random subject.
# They are almost the same, and anyway, we only need this for visualization purposes.
for subject in dataset.cameras():
if 'orientation' in dataset.cameras()[subject][args.viz_camera]:
rot = dataset.cameras()[
subject][args.viz_camera]['orientation']
break
prediction = camera_to_world(prediction, R=rot, t=0)
# We don't have the trajectory, but at least we can rebase the height
dataset = CustomDataset(custom_dataset)
print('Loading 2D detections...')
keypoints = np.load(custom_dataset, allow_pickle=True)
keypoints_metadata = keypoints['metadata'].item()
keypoints_symmetry = keypoints_metadata['keypoints_symmetry']
kps_left, kps_right = list(keypoints_symmetry[0]), list(keypoints_symmetry[1])
joints_left, joints_right = list(dataset.skeleton().joints_left()), list(
dataset.skeleton().joints_right())
keypoints = keypoints['positions_2d'].item()
for subject in keypoints.keys():
for action in keypoints[subject]:
for cam_idx, kps in enumerate(keypoints[subject][action]):
# Normalize camera frame
cam = dataset.cameras()[subject][cam_idx]
kps[..., :2] = normalize_screen_coordinates(kps[..., :2],
w=cam['res_w'],
h=cam['res_h'])
keypoints[subject][action][cam_idx] = kps
def fetch(subjects, action_filter=None, subset=1, parse_3d_poses=True):
out_poses_3d = []
out_poses_2d = []
out_camera_params = []
for subject in subjects:
for action in keypoints[subject].keys():
if action_filter is not None:
found = False
for a in action_filter:
class Predictor:
def __init__(self,
dataset_path,
checkpoint_path,
input_video_path=None,
export_path=None,
output_path=None,
with_cude=False):
self.with_cuda = with_cude
self.dataset_path = dataset_path
self.export_path = export_path
self.output_path = output_path
self.input_video_path = input_video_path
self.dataset = CustomDataset(self.dataset_path)
self.keypoints = None
self.keypoints_left = None
self.keypoints_right = None
self.joints_left = None
self.joints_right = None
self.checkpoint = torch.load(checkpoint_path,
map_location=lambda storage, loc: storage)
self.model = None
self.init_keypoints()
self.valid_poses = self.keypoints["detectron2"]["custom"]
self.init_model()
self.test_generator = None
self.init_generator()
self.prediction = None
self.make_prediction()
def export_prediction(self):
if self.export_path is not None:
np.save(self.export_path, self.prediction)
def init_model(self):
self.model = TemporalModel(self.valid_poses[0].shape[-2],
self.valid_poses[0].shape[-1],
self.dataset.skeleton().num_joints(),
filter_widths=[3, 3, 3, 3, 3],
causal=False,
dropout=0.25,
channels=1024,
dense=False)
self.model.load_state_dict(self.checkpoint['model_pos'])
def init_keypoints(self):
self.keypoints = np.load(self.dataset_path, allow_pickle=True)
keypoints_metadata = self.keypoints['metadata'].item()
keypoints_symmetry = keypoints_metadata['keypoints_symmetry']
self.keypoints_left, self.keypoints_right = list(
keypoints_symmetry[0]), list(keypoints_symmetry[1])
self.joints_left, self.joints_right = list(
self.dataset.skeleton().joints_left()), list(
self.dataset.skeleton().joints_right())
self.keypoints = self.keypoints['positions_2d'].item()
for subject in self.keypoints.keys():
for action in self.keypoints[subject]:
for cam_idx, kps in enumerate(self.keypoints[subject][action]):
# Normalize camera frame
cam = self.dataset.cameras()[subject][cam_idx]
kps[..., :2] = normalize_screen_coordinates(kps[..., :2],
w=cam['res_w'],
h=cam['res_h'])
self.keypoints[subject][action][cam_idx] = kps
def init_generator(self):
receptive_field = self.model.receptive_field()
pad = (receptive_field - 1) // 2
causal_shift = 0
self.test_generator = UnchunkedGenerator(
None,
None,
self.valid_poses,
pad=pad,
causal_shift=causal_shift,
augment=False,
kps_left=self.keypoints_left,
kps_right=self.keypoints_right,
joints_left=self.joints_left,
joints_right=self.joints_right)
def make_prediction(self):
if self.with_cuda:
self.model = self.model.cuda()
with torch.no_grad():
self.model.eval()
for _, batch, batch_2d in self.test_generator.next_epoch():
inputs_2d = torch.from_numpy(batch_2d.astype('float32'))
if self.with_cuda:
inputs_2d = inputs_2d.cuda()
predicted_3d_pos = self.model(inputs_2d)
if self.test_generator.augment_enabled():
predicted_3d_pos[1, :, :, 0] *= -1
predicted_3d_pos[1, :, self.joints_left +
self.joints_right] = predicted_3d_pos[
1, :,
self.joints_right + self.joints_left]
predicted_3d_pos = torch.mean(predicted_3d_pos,
dim=0,
keepdim=True)
predicted_3d_pos = predicted_3d_pos.squeeze(0).cpu().numpy()
rot = self.dataset.cameras()['detectron2'][0]['orientation']
predicted_3d_pos = camera_to_world(predicted_3d_pos, R=rot, t=0)
predicted_3d_pos[:, :, 2] -= np.min(predicted_3d_pos[:, :, 2])
self.prediction = predicted_3d_pos
def plot_pose(self, pose_index=0):
pose = make_pose(self.prediction.tolist()[pose_index])
pose.prepare_plot()
pose.plot()
cam_idx].shape[0]
assert keypoints[subject][action][cam_idx].shape[0] >= mocap_length
if keypoints[subject][action][cam_idx].shape[0] > mocap_length:
# Shorten sequence
keypoints[subject][action][cam_idx] = keypoints[subject][
action][cam_idx][:mocap_length]
assert len(keypoints[subject][action]) == len(
dataset[subject][action]['positions_3d'])
for subject in keypoints.keys():
for action in keypoints[subject]:
for cam_idx, kps in enumerate(keypoints[subject][action]):
# Normalize camera frame
cam = dataset.cameras()[subject][cam_idx]
kps[..., :2] = normalize_screen_coordinates(kps[..., :2],
w=cam['res_w'],
h=cam['res_h'])
keypoints[subject][action][cam_idx] = kps
# # print(keypoints)
# for k, v in keypoints.items():
# print(k,v['custom'][0].shape)
# # print(v['custom'][0][0][0])
# subjects_train = args.subjects_train.split(',')
# subjects_semi = [] if not args.subjects_unlabeled else args.subjects_unlabeled.split(',')
# if not args.render:
# subjects_test = args.subjects_test.split(',')
# else: