def binocular_2d_proj_generator(u, consts):
"""Generate binocular 2D joint position projections.
Generates bone lengths and joint angles from respective models then uses
skeleton definition to convert to 3D joint positions in global coordinate
system, before projecting to two sets of 2D image coordinates using two
offset generated camera models.
"""
n_batch = u.shape[0]
input_sizes = [
consts['n_bone_length_input'] + consts['n_joint_angle_input'],
consts['n_camera_input']
]
u_ske, u_cam, = partition(u, input_sizes)
joint_pos_3d = joint_3d_pos_generator(u_ske, consts)
cam_foc, cam_pos, cam_ang = camera_generator(u_cam, consts)
cam_mtx_1 = theano_renderer.camera_matrix_batch(
cam_foc, cam_pos + consts['cam_pos_offset'],
cam_ang + consts['cam_ang_offset'])
cam_mtx_2 = theano_renderer.camera_matrix_batch(
cam_foc, cam_pos - consts['cam_pos_offset'],
cam_ang - consts['cam_ang_offset'])
joint_pos_2d_hom_1 = tt.batched_dot(cam_mtx_1, joint_pos_3d)
joint_pos_2d_1 = (joint_pos_2d_hom_1[:, :2] /
joint_pos_2d_hom_1[:, 2][:, None, :])
joint_pos_2d_hom_2 = tt.batched_dot(cam_mtx_2, joint_pos_3d)
joint_pos_2d_2 = (joint_pos_2d_hom_2[:, :2] /
joint_pos_2d_hom_2[:, 2][:, None, :])
return tt.concatenate([
joint_pos_2d_1.reshape((n_batch, -1)),
joint_pos_2d_2.reshape((n_batch, -1))
], 1)
def joint_3d_pos_var_root_generator(u, consts):
"""Generate 3D joint positions with variable root position.
Generates bone lengths and joint angles from respective models then uses
skeleton definition to convert to 3D joint positions in global coordinate
system.
"""
input_sizes = [
consts['n_bone_length_input'], consts['n_joint_angle_input'], 3
]
u_len, u_ang, u_pos = partition(u, input_sizes)
bone_lengths = bone_lengths_generator(u_len, consts)
joint_angles = joint_angles_generator(u_ang, consts)
root_position = root_position_generator(u_pos, consts)
joint_pos_3d = tt.stack(
theano_renderer.joint_positions_batch(
consts['skeleton'],
joint_angles,
consts['fixed_joint_angles'],
lengths=bone_lengths,
lengths_map=consts['bone_lengths_map'],
skip=consts['joints_to_skip']), 2)
joint_pos_3d = tt.inc_subtensor(joint_pos_3d[:, :3, :],
root_position[:, :, None])
return joint_pos_3d
def joint_angles_generator(u, consts):
"""Generate joint angles from VAE decoder model."""
h, n = partition(
u, [consts['n_joint_angle_latent'], consts['n_joint_angle'] * 2])
ang_cos_sin_mean, ang_cos_sin_std = joint_angles_cos_sin_vae_decoder(
h, consts['joint_angles_vae_decoder_layers'], consts['n_joint_angle'])
ang_cos_sin = ang_cos_sin_mean + ang_cos_sin_std * n
return tt.arctan2(ang_cos_sin[:, consts['n_joint_angle']:],
ang_cos_sin[:, :consts['n_joint_angle']])
def energy_func_hier_binocular(state, y_data, consts):
state_partition = [
consts['n_bone_length_input'], consts['n_joint_angle_latent'],
consts['n_joint_angle'] * 2, 1, 1, 1
]
(log_bone_lengths, joint_ang_latent, joint_ang_cos_sin, cam_pos_x,
cam_pos_y, log_cam_pos_z) = partition(state, state_partition)
ang_cos_sin_mean, ang_cos_sin_std = joint_angles_cos_sin_vae_decoder(
joint_ang_latent[None, :], consts['joint_angles_vae_decoder_layers'],
consts['n_joint_angle'])
joint_angles = tt.arctan2(joint_ang_cos_sin[consts['n_joint_angle']:],
joint_ang_cos_sin[:consts['n_joint_angle']])
bone_lengths = tt.exp(log_bone_lengths)
joint_pos_3d = tt.stack(
theano_renderer.joint_positions(consts['skeleton'],
joint_angles,
consts['fixed_joint_angles'],
lengths=bone_lengths,
lengths_map=consts['bone_lengths_map'],
skip=consts['joints_to_skip']), 1)
cam_foc = consts['cam_foc']
cam_pos = tt.concatenate([cam_pos_x, cam_pos_y, tt.exp(log_cam_pos_z)])
cam_ang = consts['cam_ang']
cam_mtx_1 = theano_renderer.camera_matrix(
cam_foc, cam_pos + consts['cam_pos_offset'],
cam_ang + consts['cam_ang_offset'])
cam_mtx_2 = theano_renderer.camera_matrix(
cam_foc, cam_pos - consts['cam_pos_offset'],
cam_ang - consts['cam_ang_offset'])
joint_pos_2d_hom_1 = tt.dot(cam_mtx_1, joint_pos_3d)
joint_pos_2d_1 = joint_pos_2d_hom_1[:2] / joint_pos_2d_hom_1[2]
joint_pos_2d_hom_2 = tt.dot(cam_mtx_2, joint_pos_3d)
joint_pos_2d_2 = joint_pos_2d_hom_2[:2] / joint_pos_2d_hom_2[2]
y_model = tt.concatenate(
[joint_pos_2d_1.flatten(),
joint_pos_2d_2.flatten()], 0)
log_lengths_minus_mean = log_bone_lengths - consts['log_lengths_mean']
return 0.5 * (
(y_data - y_model).dot(y_data - y_model) /
consts['output_noise_std']**2 +
(((joint_ang_cos_sin - ang_cos_sin_mean) / ang_cos_sin_std)**2).sum() +
joint_ang_latent.dot(joint_ang_latent) + log_lengths_minus_mean.dot(
sla.solve_upper_triangular(
consts['log_lengths_covar_chol'],
sla.solve_lower_triangular(consts['log_lengths_covar_chol'].T,
log_lengths_minus_mean))) +
((cam_pos_x - consts['cam_pos_x_mean']) / consts['cam_pos_x_std'])**2 +
((cam_pos_y - consts['cam_pos_y_mean']) / consts['cam_pos_y_std'])**2 +
((log_cam_pos_z - consts['log_cam_pos_z_mean']) /
consts['log_cam_pos_z_std'])**2)[0]
def population_seq_generator(u, consts):
"""Generate simulated Lotka-Volterra model population sequences.
If a vector of random inputs `u` is provided, a single (pair) of sequences
is returned. If a matrix of random inputs `u` is provided, multiple
simulations are run in parallel as a batch, with it assumed the first
dimension of `u` corresponds to the batch dimension.
Args:
u: Tensor variable representing random inputs to generator. Either
a one-dimensional vector or two-dimensional vector for batch
generation (with first axis batch dimension).
consts: Dictionary of fixed generator constants. Should define
dt: integrator time-step,
y1_init: prey population initial state,
y2_init: predator population initial state,
log_z_std: standard dev. of log-normal prior on model parameters,
log_z_mean: mean of log-normal prior on model parameters.
Returns:
Tensor variable representing generated population sequences. If a
vector of random inputs is provided the output will be a vector of
interleaved y1 and y2 sequences [y1(1), y2(1) ... y1(T), y2(T)]. If
a matrix of random inputs this will be matrix with each row
corresponding to an interleaving of y1 and y2 sequences for the
corresponding simulation in the batch.
"""
u_z, u_n = partition(u, [consts['n_param'], consts['n_step'] * 2])
# Sample model parameters from log-normal prior
z = param_generator(u_z, consts)
# Extract initial population states and replicate to batch size
y1_init = tt.tile(
tt.constant(consts['y1_init'], 'y1_init', 0, th.config.floatX),
u.shape[0])
y2_init = tt.tile(
tt.constant(consts['y2_init'], 'y2_init', 0, th.config.floatX),
u.shape[0])
# Extract realisations of SDE white noise process terms as noise sequences
n1_seq = u_n[:, ::2] * consts['noise_std']
n2_seq = u_n[:, 1::2] * consts['noise_std']
# Iterate integrator steps given noise sequences and model parameters
[y1_seq, y2_seq], updates = th.scan(
fn=lotka_volterra_step,
sequences=[n1_seq.T, n2_seq.T],
outputs_info=[y1_init, y2_init],
non_sequences=[z.T, consts['dt']],
)
# Stack y1 and y2 sequences together to produce generator output
return tt.stack([y1_seq.T, y2_seq.T], axis=2).reshape([u.shape[0], -1])
def noisy_monocular_2d_proj_generator(u, consts):
"""Generate noisy monocular 2D joint position projections.
Generates bone lengths and joint angles from respective models then uses
skeleton definition to convert to 3D joint positions in global coordinate
system, before projecting to 2D image coordinates using a single generated
camera model and adding Gaussian observation noise to projections.
"""
input_sizes = [
consts['n_bone_length_input'] + consts['n_joint_angle_input'] +
consts['n_camera_input'], consts['n_joint'] * 2
]
u_pos, u_noi = partition(u, input_sizes)
return (monocular_2d_proj_generator(u_pos, consts) +
consts['output_noise_std'] * u_noi)
def joint_3d_pos_generator_hier(state, consts):
state_partition = [
consts['n_bone_length_input'],
consts['n_joint_angle_latent'],
consts['n_joint_angle'] * 2,
]
log_bone_lengths, joint_ang_latent, joint_ang_cos_sin = (partition(
state, state_partition))
joint_angles = tt.arctan2(joint_ang_cos_sin.T[consts['n_joint_angle']:].T,
joint_ang_cos_sin.T[:consts['n_joint_angle']].T)
bone_lengths = tt.exp(log_bone_lengths)
return tt.squeeze(
tt.stack(
theano_renderer.joint_positions_batch(
consts['skeleton'],
joint_angles,
consts['fixed_joint_angles'],
lengths=bone_lengths,
lengths_map=consts['bone_lengths_map'],
skip=consts['joints_to_skip']), 2))
def joint_3d_pos_generator(u, consts):
"""Generate 3D joint positions.
Generates bone lengths and joint angles from respective models then uses
skeleton definition to convert to 3D joint positions in global coordinate
system.
"""
input_sizes = [
consts['n_bone_length_input'], consts['n_joint_angle_input']
]
u_len, u_ang = partition(u, input_sizes)
bone_lengths = bone_lengths_generator(u_len, consts)
joint_angles = joint_angles_generator(u_ang, consts)
return tt.stack(
theano_renderer.joint_positions_batch(
consts['skeleton'],
joint_angles,
consts['fixed_joint_angles'],
lengths=bone_lengths,
lengths_map=consts['bone_lengths_map'],
skip=consts['joints_to_skip']), 2)
def inputs_to_state(u, consts):
input_sizes = [
consts['n_bone_length_input'], consts['n_joint_angle_latent'],
consts['n_joint_angle'] * 2, consts['n_camera_input']
]
u_len, joint_ang_latent, joint_ang_noise, u_cam = partition(u, input_sizes)
log_bone_lengths = (consts['log_lengths_mean'] +
u_len.dot(consts['log_lengths_covar_chol']))
joint_ang_cos_sin_mean, joint_ang_cos_sin_log_var = (
consts['joint_angles_cos_sin_vae'].x_gvn_z(joint_ang_latent))
joint_ang_cos_sin = tt.squeeze(joint_ang_cos_sin_mean +
tt.exp(0.5 * joint_ang_cos_sin_log_var) *
joint_ang_noise)
cam_pos_x = consts['cam_pos_x_mean'] + consts['cam_pos_x_std'] * u_cam[0:1]
cam_pos_y = consts['cam_pos_y_mean'] + consts['cam_pos_y_std'] * u_cam[1:2]
log_cam_pos_z = (consts['log_cam_pos_z_mean'] +
consts['log_cam_pos_z_std'] * u_cam[2:3])
return tt.concatenate([
log_bone_lengths, joint_ang_latent, joint_ang_cos_sin, cam_pos_x,
cam_pos_y, log_cam_pos_z
], 0)
def monocular_2d_proj_generator(u, consts):
"""Generate monocular 2D joint position projections.
Generates bone lengths and joint angles from respective models then uses
skeleton definition to convert to 3D joint positions in global coordinate
system, before projecting to 2D image coordinates using a single generated
camera model.
"""
input_sizes = [
consts['n_bone_length_input'] + consts['n_joint_angle_input'],
consts['n_camera_input']
]
u_ske, u_cam = partition(u, input_sizes)
joint_pos_3d = joint_3d_pos_generator(u_ske, consts)
cam_foc, cam_pos, cam_ang = camera_generator(u_cam, consts)
camera_matrix = theano_renderer.camera_matrix_batch(
cam_foc, cam_pos, cam_ang)
joint_pos_2d_hom = tt.batched_dot(camera_matrix, joint_pos_3d)
joint_pos_2d = (joint_pos_2d_hom[:, :2] /
joint_pos_2d_hom[:, 2][:, None, :])
return joint_pos_2d
def mnist_region_generator(u, consts):
"""Generate region of MNIST digit images (as flat vectors)."""
h, n = partition(u, [consts['n_latent'], consts['n_observed']])
m = mnist_vae_decoder(h, consts['mnist_vae_decoder_layers'])
m_obs = m[:, consts['observed_slice']]
return m_obs + consts['output_std'] * n
def mnist_generator(u, consts):
"""Generate MNIST digit images (as flat vectors)."""
h, n = partition(u, [consts['n_latent'], consts['n_pixel']])
m = mnist_vae_decoder(h, consts['mnist_vae_decoder_layers'])
return m + consts['output_std'] * n