From 8016f0f945be36bdec4f3300bb37a8d34efc2706 Mon Sep 17 00:00:00 2001
From: Piv <18462828+Piv200@users.noreply.github.com>
Date: Sun, 8 Aug 2021 22:11:50 +0930
Subject: [PATCH] Add coordinates generation implementation

---
 unsupervised/warp.py       | 27 +++++++++++++++++++++++++--
 unsupervised/warp_tests.py | 18 ++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/unsupervised/warp.py b/unsupervised/warp.py
index 28712e9..190aa4d 100644
--- a/unsupervised/warp.py
+++ b/unsupervised/warp.py
@@ -72,7 +72,28 @@ def pose_vec2mat(vec):
     return transform_mat
 
 
-def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics):
+def image_coordinate(batch, height, width):
+    """
+    Construct a tensor for the given height/width with elements the homogenous coordinates for the pixel
+    :param batch: Number of images in a batch
+    :param height: Height of image
+    :param width: Width of image
+    :return: Tensor of shape (B, height, width, 3), homogenous coordinates for an image.
+        Coordinates are in order [x, y, 1]
+    """
+    x_coords = tf.range(width)
+    y_coords = tf.range(height)
+
+    x_mesh, y_mesh = tf.meshgrid(x_coords, y_coords)
+
+    ones_mesh = tf.cast(tf.ones([height, width]), tf.int32)
+
+    stacked = tf.stack([x_mesh, y_mesh, ones_mesh], axis=2)
+
+    return tf.repeat(tf.expand_dims(stacked, axis=0), batch, axis=0)
+
+
+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
 
@@ -88,12 +109,14 @@ def projective_inverse_warp(target_img, source_img, depth, pose, intrinsics):
     :param depth: Tensor, (batch, height, width, 1)
     :param pose: (batch, 6)
     :param intrinsics: (batch, 3, 3)
+    :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)
 
     # Convert intrinsics matrix (3x3) to (4x4) so it can be multiplied by the pose net
-    intrinsics_4x4 =
+    # intrinsics_4x4 =
 
     # Calculate inverse of the 4x4 intrinsics matrix
     tf.linalg.inv()
diff --git a/unsupervised/warp_tests.py b/unsupervised/warp_tests.py
index d21f321..a96c4ca 100644
--- a/unsupervised/warp_tests.py
+++ b/unsupervised/warp_tests.py
@@ -21,6 +21,24 @@ class MyTestCase(unittest.TestCase):
 
         # TODO: Element-wise checks...
 
+    def test_coordinates(self):
+        height = 1000
+        width = 2000
+        coords = warp.image_coordinate(8, height, width)
+
+        self.assertEqual(coords.shape, [8, height, width, 3])
+        self.assertEqual(coords[0, 0, 0, 0], 0)
+        self.assertEqual(coords[0, 0, 0, 1], 0)
+        self.assertEqual(coords[0, 0, 0, 2], 1)
+
+        self.assertEqual(coords[0, height - 1, 0, 0], 0)
+        self.assertEqual(coords[0, height - 1, 0, 1], height - 1)
+        self.assertEqual(coords[0, height - 1, 0, 2], 1)
+
+        self.assertEqual(coords[0, height - 1, width - 1, 0], width - 1)
+        self.assertEqual(coords[0, height - 1, width - 1, 1], height - 1)
+        self.assertEqual(coords[0, height - 1, width - 1, 2], 1)
+
 
 if __name__ == '__main__':
     unittest.main()