def spherical_expectation(spherical_probabilities):
"""Compute the expectation (a vector) from normalized spherical distribtuions.
We define the spherical expectation as the integral of r*P(r)*dr where r is
a unit direction vector in 2-sphere. We compute the discretized version on a
spherical equirectangular map. To correctly use this function, the input has
to be normalized properly using spherical_normalization().
Args:
spherical_probabilities: [BATCH, HEIGHT, WIDTH, N] spherical distributions
in equirectangular form.
Returns:
expectation [BATCH, N, 3]
"""
shape = spherical_probabilities.shape.as_list()
height, width, channels = shape[1], shape[2], shape[3]
spherical = tf.expand_dims(
geometry.generate_equirectangular_grid([height, width]), 0)
unit_directions = geometry.spherical_to_cartesian(spherical)
axis_convert = tf.constant([[1., 0., 0.], [0., 0., -1.], [0., 1., 0.]])
unit_directions = tf.squeeze(
tf.matmul(axis_convert,
tf.expand_dims(unit_directions, -1),
transpose_a=True), -1)
unit_directions = tf.tile(tf.expand_dims(unit_directions, -2),
[1, 1, 1, channels, 1])
weighted = spherical_probabilities * equirectangular_area_weights(height)
expectation = tf.reduce_sum(unit_directions * tf.expand_dims(weighted, -1),
[1, 2])
return expectation
def von_mises_fisher(mean, concentration, shape):
"""Generate von Mises-Fisher distribution on spheres.
This function samples probabilities from tensorflow_probability.VonMisesFisher
on equirectangular grids of a sphere. The height dimension of the output
ranges from pi/2 (top) to -pi/2 (bottom). The width dimension ranges from
0 (left) to 2*pi (right).
Args:
mean: [BATCH, N, 3] a float tensor representing the unit direction of
the mean.
concentration: (float) a measure of concentration (a reciprocal measure of
dispersion, so 1/kappa is analogous to variance). concentration=0
indicates a uniform distribution over the unit sphere,
and concentration=+inf indicates a delta function at the mean direction.
shape: a 2-d list represents the dimension (height, width) of the output.
Returns:
A 4-D tensor [BATCH, HEIGHT, WIDTH, N] represents the raw probabilities
of the distribution. (surface integral != 1)
Raises:
ValueError: Input argument 'shape' is not valid.
ValueError: Input argument 'mean' has wrong dimensions.
"""
with tf.name_scope(None, 'von_mises_fisher', [mean, concentration, shape]):
if not isinstance(shape, list) or len(shape) != 2:
raise ValueError("Input argument 'shape' is not valid.")
if mean.shape[-1] != 3:
raise ValueError("Input argument 'mean' has wrong dimensions.")
batch, channels = mean.shape[0], mean.shape[1]
height, width = shape[0], shape[1]
spherical_grid = geometry.generate_equirectangular_grid(shape)
cartesian = geometry.spherical_to_cartesian(spherical_grid)
axis_convert = tf.constant([[1., 0., 0.], [0., 0., -1.], [0., 1., 0.]])
cartesian = tf.squeeze(
tf.matmul(axis_convert,
tf.expand_dims(cartesian, -1),
transpose_a=True), -1)
cartesian = tf.tile(cartesian[tf.newaxis, tf.newaxis, :],
[batch, channels, 1, 1, 1])
mean = tf.tile(mean[:, :, tf.newaxis, tf.newaxis],
[1, 1, height, width, 1])
vmf = tfp.distributions.VonMisesFisher(mean_direction=mean,
concentration=[concentration])
spherical_gaussian = vmf.prob(cartesian)
return tf.transpose(spherical_gaussian, [0, 2, 3, 1])
def rotate_image_on_pano(images, rotations, fov, output_shape):
"""Transform perspective images to equirectangular images after rotations.
Return equirectangular panoramic images in which the input perspective images
embedded in after the rotation R from the input images' frame to the target
frame. The image with the field of view "fov" centered at camera's look-at -Z
axis is projected onto the pano. The -Z axis corresponds to the spherical
coordinates (pi/2, pi/2) which is (HEIGHT/2, WIDTH/4) on the pano.
Args:
images: [BATCH, HEIGHT, WIDTH, CHANNEL] perspective view images.
rotations: [BATCH, 3, 3] rotations matrices.
fov: (float) images' field of view in degrees.
output_shape: a 2-D list of output dimension [height, width].
Returns:
equirectangular images [BATCH, height, width, CHANNELS].
"""
with tf.name_scope(None, 'rotate_image_on_pano',
[images, rotations, fov, output_shape]):
if len(images.shape) != 4:
raise ValueError("'images' has the wrong dimensions.")
if rotations.shape[-2:] != [3, 3]:
raise ValueError("'rotations' must have 3x3 dimensions.")
shape = images.shape.as_list()
batch, height, width = shape[0], shape[1], shape[2]
# Generate a mesh grid on a sphere.
spherical = geometry.generate_equirectangular_grid(output_shape)
cartesian = geometry.spherical_to_cartesian(spherical)
cartesian = tf.tile(cartesian[tf.newaxis, :, :, :, tf.newaxis],
[batch, 1, 1, 1, 1])
axis_convert = tf.constant([[0., -1., 0.], [0., 0., 1.], [1., 0., 0.]])
cartesian = tf.matmul(axis_convert, cartesian)
cartesian = tf.squeeze(
tf.matmul(rotations[:, tf.newaxis, tf.newaxis], cartesian), -1)
# Only take one hemisphere. (camera lookat direction)
hemisphere_mask = tf.cast(cartesian[:, :, :, -1:] < 0, tf.float32)
image_coordinates = cartesian[:, :, :, :2] / cartesian[:, :, :, -1:]
x, y = tf.split(image_coordinates, [1, 1], -1)
# Map pixels on equirectangular pano to perspective image.
nx = -x * width / (2 * tf.tan(math_utils.degrees_to_radians(
fov / 2))) + width / 2 - 0.5
ny = y * height / (2 * tf.tan(math_utils.degrees_to_radians(
fov / 2))) + height / 2 - 0.5
transformed = hemisphere_mask * tfa.image.resampler(
images, tf.concat([nx, ny], -1))
return transformed
def rotate_pano(images, rotations):
"""Rotate Panoramic images.
Convert the spherical coordinates (colatitude, azimuth) to Cartesian (x, y, z)
then apply SO(3) rotation matrices. Finally, convert them back to spherical
coordinates and remap the equirectangular images.
Note1: The rotations are applied to the sampling sphere instead of the camera.
The camera actually rotates R^T. I_out(x) = I_in(R * x), x are points in the
camera frame.
Note2: It uses a simple linear interpolation for now instead of slerp, so the
pixel values are not accurate but visually plausible.
Args:
images: a 4-D tensor of shape `[BATCH, HEIGHT, WIDTH, CHANNELS]`.
rotations: [BATCH, 3, 3] rotation matrices.
Returns:
4-D tensor of shape `[BATCH, HEIGHT, WIDTH, CHANNELS]`.
Raises:
ValueError: if the `images` or 'rotations' has the wrong dimensions.
"""
with tf.name_scope(None, 'rotate_pano', [images, rotations]):
if len(images.shape) != 4:
raise ValueError("'images' has the wrong dimensions.")
if rotations.shape[-2:] != [3, 3]:
raise ValueError("'rotations' must have 3x3 dimensions.")
shape = images.shape.as_list()
batch, height, width = shape[0], shape[1], shape[2]
spherical = tf.expand_dims(
geometry.generate_equirectangular_grid([height, width]), 0)
spherical = tf.tile(spherical, [batch, 1, 1, 1])
cartesian = geometry.spherical_to_cartesian(spherical)
axis_convert = tf.constant([[0., 1., 0.], [0., 0., -1.], [-1., 0.,
0.]])
cartesian = tf.matmul(axis_convert, tf.expand_dims(cartesian, -1))
rotated_cartesian = tf.matmul(rotations[:, tf.newaxis, tf.newaxis],
cartesian)
rotated_cartesian = tf.squeeze(
tf.matmul(axis_convert, rotated_cartesian, transpose_a=True), -1)
rotated_spherical = geometry.cartesian_to_spherical(rotated_cartesian)
return equirectangular_sampler(images, rotated_spherical)