def _get_angles(coords):
base_shape = [int(dim) for dim in coords.shape]
base_shape.pop(1)
base_shape[-1] = 1
coords_list = tf.unstack(coords, axis=1)
def _get_angle_for_joint(joint_idx, parent_idx, angles):
if parent_idx is None: # joint_idx should be 0
parent_bone = K.constant(
np.concatenate([
np.ones(base_shape),
np.zeros(base_shape),
np.zeros(base_shape)
],
axis=-1))
else:
parent_bone = coords_list[parent_idx] - coords_list[
joint_idx]
for child_idx in body_graph[joint_idx]:
child_bone = coords_list[child_idx] - coords_list[joint_idx]
angle = quaternion_between(parent_bone, child_bone)
angle = quaternion_to_expmap(angle)
angles.append(angle)
for child_idx in body_graph[joint_idx]:
angles = _get_angle_for_joint(child_idx, joint_idx, angles)
return angles
angles = _get_angle_for_joint(0, None, [])
return K.stack(angles, axis=1)
def _diff_to_seq(args):
diffs, start_pose = args
diffs_list = tf.unstack(diffs, axis=2)
poses = [start_pose]
for p in range(diffs.shape[2]):
poses.append(poses[p] + diffs_list[p])
return K.stack(poses, axis=2)
def _get_bone_len(arg):
bone_list = tf.unstack(arg[:, :, 0, :], axis=1)
bones = [
bone_list[j] - bone_list[i]
for i, j in zip(members_from, members_to)
]
bones = K.stack(bones, axis=1)
return K.sqrt(K.sum(K.square(bones), axis=-1) + K.epsilon())
def _get_avg_bone_len(arg):
bone_list = tf.unstack(arg[:, :, 0, :], axis=1)
bones = [
bone_list[j] - bone_list[i]
for i, j in zip(members_from, members_to)
]
bones = K.expand_dims(K.stack(bones, axis=1), axis=2)
bone_len = K.sqrt(
K.sum(K.square(bones), axis=-1, keepdims=True) + K.epsilon())
return K.mean(bone_len, axis=1, keepdims=True)
def _get_coords(args):
rotmat, bone_len = args
rotmat_list = tf.unstack(rotmat, axis=1)
bone_len_list = tf.unstack(bone_len, axis=1)
base_shape = [int(d) for d in rotmat.shape]
base_shape.pop(1)
base_shape[-2] = 1
base_shape[-1] = 1
bone_idcs = {
idx_tup: i
for i, idx_tup in enumerate(
[idx_tup for idx_tup in zip(members_from, members_to)])
}
def _get_coords_for_joint(joint_idx, parent_idx, child_angle_idx,
coords):
if parent_idx is None: # joint_idx should be 0
coords[joint_idx] = K.zeros(base_shape[:-2] + [3, 1])
parent_bone = K.constant(
np.concatenate([
np.ones(base_shape),
np.zeros(base_shape),
np.zeros(base_shape)
],
axis=-2))
else:
parent_bone = coords[parent_idx] - coords[joint_idx]
parent_bone_norm = K.sqrt(
K.sum(K.square(parent_bone), axis=-2, keepdims=True) +
K.epsilon())
parent_bone = parent_bone / parent_bone_norm
for child_idx in body_graph[joint_idx]:
child_bone = tf.matmul(rotmat_list[child_angle_idx],
parent_bone)
child_bone_idx = bone_idcs[(joint_idx, child_idx)]
child_bone = child_bone * K.reshape(
bone_len_list[child_bone_idx],
(child_bone.shape[0], 1, 1, 1))
coords[child_idx] = child_bone + coords[joint_idx]
child_angle_idx += 1
for child_idx in body_graph[joint_idx]:
child_angle_idx, coords = _get_coords_for_joint(
child_idx, joint_idx, child_angle_idx, coords)
return child_angle_idx, coords
child_angle_idx, coords = _get_coords_for_joint(0, None, 0, {})
coords = K.stack([t for i, t in sorted(coords.iteritems())],
axis=1)
coords = K.squeeze(coords, axis=-1)
return coords
def _get_rotation(arg):
coords_list = tf.unstack(arg[:, :, 0, :], axis=1)
torso_rot = tf.cross(
coords_list[left_shoulder] - coords_list[hip],
coords_list[right_shoulder] - coords_list[hip])
side_rot = K.reshape(
tf.cross(coords_list[head_top] - coords_list[hip], torso_rot),
base_shape)
theta_diff = (
(np.pi / 2) - tf.atan2(side_rot[..., 1], side_rot[..., 0])) / 2
cos_theta_diff = tf.cos(theta_diff)
sin_theta_diff = tf.sin(theta_diff)
zeros_theta = K.zeros_like(sin_theta_diff)
return K.stack(
[cos_theta_diff, zeros_theta, zeros_theta, sin_theta_diff],
axis=-1)
def call(self, inputs, **kwargs):
input_shape = K.int_shape(inputs)
tensor_input_shape = K.shape(inputs)
# Prepare broadcasting shape.
reduction_axes = list(range(len(input_shape)))
del reduction_axes[self.axis]
broadcast_shape = [1] * len(input_shape)
broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
broadcast_shape.insert(1, self.groups)
reshape_group_shape = K.shape(inputs)
group_axes = [reshape_group_shape[i] for i in range(len(input_shape))]
group_axes[self.axis] = input_shape[self.axis] // self.groups
group_axes.insert(1, self.groups)
# reshape inputs to new group shape
group_shape = [group_axes[0], self.groups] + group_axes[2:]
group_shape = K.stack(group_shape)
inputs = K.reshape(inputs, group_shape)
group_reduction_axes = list(range(len(group_axes)))
mean, variance = _moments(inputs,
group_reduction_axes[2:],
keep_dims=True)
inputs = (inputs - mean) / (K.sqrt(variance + self.epsilon))
# prepare broadcast shape
inputs = K.reshape(inputs, group_shape)
outputs = inputs
# In this case we must explicitly broadcast all parameters.
if self.scale:
broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
outputs = outputs * broadcast_gamma
if self.center:
broadcast_beta = K.reshape(self.beta, broadcast_shape)
outputs = outputs + broadcast_beta
# finally we reshape the output back to the input shape
outputs = K.reshape(outputs, tensor_input_shape)
return outputs
def _atari_state_feature_net(observation_space, net_name):
num_frames = len(observation_space.spaces)
height, width = observation_space.spaces[0].shape
input_shape = height, width, num_frames
# input state
state = Input(shape=input_shape)
# convolutional layers
conv1_32 = Conv2D(32, (8, 8), strides=(4, 4), activation='relu')
conv2_64 = Conv2D(64, (4, 4), strides=(2, 2), activation='relu')
conv3_64 = Conv2D(64, (3, 3), strides=(1, 1), activation='relu')
# if recurrent net then change input shape
if 'lstm' in net_name or 'gru' in net_name:
# recurrent net
lambda_perm_state = lambda x: K.permute_dimensions(x, [0, 3, 1, 2])
perm_state = Lambda(lambda_perm_state)(state)
dist_state = Lambda(lambda x: K.stack([x], axis=4))(perm_state)
# extract features with `TimeDistributed` wrapped convolutional layers
dist_conv1 = TimeDistributed(conv1_32)(dist_state)
dist_conv2 = TimeDistributed(conv2_64)(dist_conv1)
dist_convf = TimeDistributed(conv3_64)(dist_conv2)
feature = TimeDistributed(Flatten())(dist_convf)
# specify net type for the following layer
if 'lstm' in net_name:
net = LSTM
elif 'gru' in net_name:
net = GRU
elif 'fc' in net_name:
# fully connected net
# extract features with convolutional layers
conv1 = conv1_32(state)
conv2 = conv2_64(conv1)
convf = conv3_64(conv2)
feature = Flatten()(convf)
# specify net type for the following layer
net = Dense
else:
raise ValueError('`net_name` is not recognized')
return state, feature, net
def _get_angles_mask(coord_masks):
base_shape = [int(dim) for dim in coord_masks.shape]
base_shape.pop(1)
base_shape[-1] = 1
coord_masks_list = tf.unstack(coord_masks, axis=1)
def _get_angle_mask_for_joint(joint_idx, angles_mask):
for child_idx in body_graph[joint_idx]:
angles_mask.append(coord_masks_list[child_idx]
) # * coord_masks_list[joint_idx]
for child_idx in body_graph[joint_idx]:
angles_mask = _get_angle_mask_for_joint(
child_idx, angles_mask)
return angles_mask
angles_mask = _get_angle_mask_for_joint(0, [])
return K.stack(angles_mask, axis=1)