Start adding pose warp conversions

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

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()