More warp implementation

unsupervised/warp.py

@@ -44,9 +44,8 @@ def pose_vec2mat(vec):
     Args:
         vec: 6DoF parameters in the order of tx, ty, tz, rx, ry, rz -- [B, 6]
     Returns:
-        A transformation matrix -- [B, 4, 4]
+        A transformation matrix -- [B, 3, 4]
     """
-    # TODO: FIXME
     batch_size, _ = vec.get_shape().as_list()
     translation = tf.slice(vec, [0, 0], [-1, 3])
     translation = tf.expand_dims(translation, -1)
@@ -55,10 +54,7 @@ def pose_vec2mat(vec):
     rz = tf.slice(vec, [0, 5], [-1, 1])
     rot_mat = euler_to_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
 
 
@@ -83,15 +79,6 @@ def image_coordinate(batch, height, width):
     return tf.repeat(tf.expand_dims(stacked, axis=0), batch, axis=0)
 
 
-def intrinsics_vector_to_matrix(intrinsics):
-    """
-    Convert 4 element
-    :param intrinsics: Tensor of shape (B, 4), intrinsics for each image
-    :return: Tensor of shape (B, 4, 4), intrinsics for each batch
-    """
-    pass
-
-
 def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics, coordinates):
     """
     Calculate the reprojected image from the source to the target, based on the given depth, pose and intrinsics
@@ -109,27 +96,35 @@ def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics, coo
     :param source_img: Tensor, same shape as target_img
     :param depth: Tensor, (batch, height, width, 1)
     :param pose: (batch, 6)
-    :param intrinsics: (batch, 4) (fx, fy, px, py) TODO: Intrinsics per image (per source/target image)?
+    :param intrinsics: (batch, 3, 3) TODO: Intrinsics per image (per source/target image)?
     :param coordinates: (batch, height, width, 3) - coordinates for the image. Pass this in so it doesn't need to be
         calculated on every warp step
     :return: The source image reprojected to the target
     """
     # Convert pose vector (output of pose net) to pose matrix (4x4)
-    pose_4x4 = pose_vec2mat(pose)
+    pose_3x4 = pose_vec2mat(pose)
 
     # 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()
+    intrinsics_inverse = tf.linalg.inv(intrinsics)
 
-    # Create grid (or array?) of homogenous coordinates
+    # Create grid of homogenous coordinates
     grid_coords = image_coordinate(*depth.shape)
     # Flatten the image coords to [B, 3, height * width] so each point can be used in calculations
     grid_coords = tf.transpose(tf.reshape(grid_coords, [0, 2, 1]))
 
-    # Get grid coordinates as array
+    # TODO: Do we need to transpose?
+    depth_flat = tf.transpose(tf.reshape(depth, [0, 2, 1]))
 
     # Do the function
+    sample_coordinates = tf.matmul(tf.matmul(intrinsics, pose_3x4),
+                                   tf.concat([depth_flat * tf.matmul(intrinsics_inverse, grid_coords),
+                                              tf.ones(depth_flat.shape)], axis=1))
+
+    # Normalise the x/y axes (divide by z axis)
+
+    # Reshape back to image coordinates
 
     # sample from the source image using the coordinates applied by the function