diff --git a/unsupervised/warp.py b/unsupervised/warp.py index 7cc19b3..28712e9 100644 --- a/unsupervised/warp.py +++ b/unsupervised/warp.py @@ -1,3 +1,77 @@ +import numpy as np +import tensorflow as tf + + +def euler_to_rotation_matrix(x, y, z): + """ + + :param x: Tensor of shape (B, 1) - x axis rotation + :param y: Tensor of shape (B, 1) - y axis rotation + :param z: Tensor of shape (B, 1) - z axis rotation + :return: Rotation matrix for the given euler anglers, in the order rotation(x).rotation(y).rotation(z) + """ + B = tf.shape(z)[0] + + # Euler angles should be between -pi and pi, clip so the pose network is coerced to this range + z = tf.clip_by_value(z, -np.pi, np.pi) + y = tf.clip_by_value(y, -np.pi, np.pi) + x = tf.clip_by_value(x, -np.pi, np.pi) + + # Expand to B x 1 x 1 + z = tf.expand_dims(tf.expand_dims(z, -1), -1) + y = tf.expand_dims(tf.expand_dims(y, -1), -1) + x = tf.expand_dims(tf.expand_dims(x, -1), -1) + + zeros = tf.zeros([B, 1, 1]) + ones = tf.ones([B, 1, 1]) + + cosx = tf.cos(x) + sinx = tf.sin(x) + rotx_1 = tf.concat([ones, zeros, zeros], axis=3) + rotx_2 = tf.concat([zeros, cosx, -sinx], axis=3) + rotx_3 = tf.concat([zeros, sinx, cosx], axis=3) + xmat = tf.concat([rotx_1, rotx_2, rotx_3], axis=2) + + cosz = tf.cos(z) + sinz = tf.sin(z) + rotz_1 = tf.concat([cosz, -sinz, zeros], axis=3) + rotz_2 = tf.concat([sinz, cosz, zeros], axis=3) + rotz_3 = tf.concat([zeros, zeros, ones], axis=3) + zmat = tf.concat([rotz_1, rotz_2, rotz_3], axis=2) + + cosy = tf.cos(y) + siny = tf.sin(y) + roty_1 = tf.concat([cosy, zeros, siny], axis=3) + roty_2 = tf.concat([zeros, ones, zeros], axis=3) + roty_3 = tf.concat([-siny, zeros, cosy], axis=3) + ymat = tf.concat([roty_1, roty_2, roty_3], axis=2) + + rotMat = tf.matmul(tf.matmul(xmat, ymat), zmat) + return rotMat + + +def pose_vec2mat(vec): + """Converts 6DoF parameters to transformation matrix + Args: + vec: 6DoF parameters in the order of tx, ty, tz, rx, ry, rz -- [B, 6] + Returns: + A transformation matrix -- [B, 4, 4] + """ + batch_size, _ = vec.get_shape().as_list() + translation = tf.slice(vec, [0, 0], [-1, 3]) + translation = tf.expand_dims(translation, -1) + rx = tf.slice(vec, [0, 3], [-1, 1]) + ry = tf.slice(vec, [0, 4], [-1, 1]) + rz = tf.slice(vec, [0, 5], [-1, 1]) + rot_mat = euler_to_rotation_matrix(rx, ry, rz) + rot_mat = tf.squeeze(rot_mat, axis=[1]) + filler = tf.constant([0.0, 0.0, 0.0, 1.0], shape=[1, 1, 4]) + filler = tf.tile(filler, [batch_size, 1, 1]) + transform_mat = tf.concat([rot_mat, translation], axis=2) + transform_mat = tf.concat([transform_mat, filler], axis=1) + return transform_mat + + def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics): """ Calculate the reprojected image from the source to the target, based on the given depth, pose and intrinsics @@ -12,8 +86,19 @@ def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics): :param target_img: Tensor (batch, height, width, 3) :param source_img: Tensor, same shape as target_img :param depth: Tensor, (batch, height, width, 1) - :param pose: (batch, 3, 3) + :param pose: (batch, 6) :param intrinsics: (batch, 3, 3) :return: The source image reprojected to the target """ + # Convert pose vector (output of pose net) to pose matrix (4x4) + + # Convert intrinsics matrix (3x3) to (4x4) so it can be multiplied by the pose net + intrinsics_4x4 = + + # Calculate inverse of the 4x4 intrinsics matrix + tf.linalg.inv() + + # Create grid of homogenous coordinates + + # pass diff --git a/unsupervised/warp_tests.py b/unsupervised/warp_tests.py new file mode 100644 index 0000000..d21f321 --- /dev/null +++ b/unsupervised/warp_tests.py @@ -0,0 +1,26 @@ +import unittest + +import numpy as np +import tensorflow as tf + +import warp + + +class MyTestCase(unittest.TestCase): + def test_euler_to_rotation_matrix(self): + # quarter rotation in every + x, y, z = tf.expand_dims(tf.expand_dims(tf.constant(np.pi / 2), 0), 0) + + # TODO: Construct expected final rotation matrix, just 3x3 using numpy, so that we can do an + # elementwise comparison later. Probably also want to check the + + rotation_matrices = warp.euler_to_rotation_matrix(x, y, z) + + self.assertEqual(rotation_matrices.shape, [1, 3, 3]) + rot_mat = rotation_matrices[0] + + # TODO: Element-wise checks... + + +if __name__ == '__main__': + unittest.main()