def rotate(input_images, rx, ry, rz):
# Create constants.
batch_size = tf.shape(input_images)[0]
height = tf.shape(input_images)[1]
width = tf.shape(input_images)[2]
batch_zero = tf.tile(tf.constant(0.0, shape = [1]), [batch_size])
batch_one = tf.tile(tf.constant(1.0, shape = [1]), [batch_size])
# Function to parse Python lists as TensorFlow matrices.
def tf_batch_matrix(matrix):
return tf.transpose(tf.stack(matrix), [2, 0, 1])
# Convert to Cartesian.
S, T = lat_long_grid([height, width])
x, y, z = lat_long_to_xyz(S, T)
X = tf.tile(tf.expand_dims(tf.reshape(tf.stack([x, y, z]), [3, height * width]), 0), [batch_size, 1, 1])
# Construct rotation matrices (for inverse warp).
rx = - rx
ry = - ry
rz = - rz
R = tf_batch_matrix([
[tf.cos(rz), - tf.sin(rz), batch_zero],
[tf.sin(rz), tf.cos(rz), batch_zero],
[batch_zero, batch_zero, batch_one]
])
R = tf.matmul(
tf_batch_matrix([
[tf.cos(ry), batch_zero, tf.sin(ry)],
[batch_zero, batch_one, batch_zero],
[- tf.sin(ry), batch_zero, tf.cos(ry)]
]),
R)
R = tf.matmul(
tf_batch_matrix([
[batch_one, batch_zero, batch_zero],
[batch_zero, tf.cos(rx), - tf.sin(rx)],
[batch_zero, tf.sin(rx), tf.cos(rx)]
]),
R)
# Rotate coordinates.
X_rotated = tf.reshape(tf.matmul(R, X), [batch_size, 3, height, width])
# Convert back to equirectangular UV.
S_rotated, T_rotated = xyz_to_lat_long(X_rotated[:, 0, :, :], X_rotated[:, 1, :, :], X_rotated[:, 2, :, :])
u, v = lat_long_to_equirectangular_uv(S_rotated, T_rotated)
return bilinear_sample(input_images, x_t = tf.zeros_like(u[0]), y_t = tf.zeros_like(v[0]), x_offset = u, y_offset = 1.0 - v)
def rectilinear_to_equirectangular(input_images, K, equirectangular_shape):
stacked_faces = stack_faces(input_images)
S, T = lat_long_grid(equirectangular_shape)
u, v = lat_long_to_rectilinear_uv(K, S, T)
return bilinear_sample(stacked_faces, u, v)
def cubic_to_equirectangular(input_images, equirectangular_shape):
stacked_faces = stack_faces(input_images)
S, T = lat_long_grid(equirectangular_shape)
u, v = lat_long_to_cube_uv(S, T)
return bilinear_sample(stacked_faces, u, v)
def project_rectilinear(input_images, K, face, face_shape):
x, y, z = rectilinear_xyz(K, face_shape, face)
S, T = xyz_to_lat_long(x, y, z)
u, v = lat_long_to_equirectangular_uv(S, T)
return bilinear_sample(input_images, u, v)
def project_face(input_images, face, cubic_shape):
x, y, z = xyz_grid(cubic_shape, face)
S, T = xyz_to_lat_long(x, y, z)
u, v = lat_long_to_equirectangular_uv(S, T)
return bilinear_sample(input_images, u, v)