Move out common functions to utils.py

PiperOrigin-RevId: 361885696

tensorflow_graphics/rendering/kernels/rasterization_backend.py

@@ -19,6 +19,7 @@
 import tensorflow as tf
 
 from tensorflow_graphics.rendering import framebuffer as fb
+from tensorflow_graphics.rendering import utils
 from tensorflow_graphics.util import shape
 
 # pylint: disable=g-import-not-at-top
@@ -38,30 +39,6 @@ class FaceCullingMode(enum.IntEnum):
   FRONT = 2
 
 
-def _transform_homogeneous(matrices, vertices):
-  """Applies 4x4 homogenous matrix transformations to xyz vertices.
-
-  The vertices are input and output as as row-major, but are interpreted as
-  column vectors multiplied on the right-hand side of the matrices. More
-  explicitly, this function computes (MV^T)^T where M represents transformation
-  matrices and V stands for vertices.
-  Since input vertices are xyz they are extended to xyzw with w=1.
-
-  Args:
-    matrices: A tensor of shape `[batch, 4, 4]` containing batches of view
-      projection matrices.
-    vertices: A tensor of shape `[batch, num_vertices, 3]` containing batches of
-      vertices, each defined by a 3D point.
-
-  Returns:
-    A [batch, N, 4] Tensor of xyzw vertices.
-  """
-  homogeneous_coord = tf.ones_like(vertices[..., 0:1])
-  vertices = tf.concat([vertices, homogeneous_coord], -1)
-
-  return tf.matmul(vertices, matrices, transpose_b=True)
-
-
 def rasterize(vertices: tf.Tensor,
               triangles: tf.Tensor,
               view_projection_matrices: tf.Tensor,
@@ -133,7 +110,7 @@ def rasterize(vertices: tf.Tensor,
     if not num_layers > 0:
       raise ValueError("num_layers must be > 0.")
 
-    vertices = _transform_homogeneous(view_projection_matrices, vertices)
+    vertices = utils.transform_homogeneous(view_projection_matrices, vertices)
     batch_size = tf.compat.dimension_value(vertices.shape[0])
 
     per_image_barycentrics = []

tensorflow_graphics/rendering/tests/utils_test.py

@@ -0,0 +1,44 @@
+# Copyright 2020 The TensorFlow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for common util functions."""
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow_graphics.rendering import utils
+from tensorflow_graphics.util import test_case
+
+
+class UtilsTest(test_case.TestCase):
+
+  @parameterized.named_parameters(
+      ('non-batched xyz', False),
+      ('batched xyz', True),
+  )
+  def test_transform_homogeneous_shapes(self, do_batched):
+    num_vertices = 10
+    batch_size = 3
+    num_channels = 3
+    vertices_shape = ([batch_size, num_vertices, num_channels]
+                      if do_batched else [num_vertices, num_channels])
+    vertices = tf.ones(vertices_shape, dtype=tf.float32)
+    matrices = tf.eye(4, dtype=tf.float32)
+    if do_batched:
+      matrices = tf.tile(matrices[tf.newaxis, ...], [batch_size, 1, 1])
+
+    transformed = utils.transform_homogeneous(matrices, vertices)
+
+    expected_shape = ([batch_size, num_vertices, 4]
+                      if do_batched else [num_vertices, 4])
+    self.assertEqual(transformed.shape, expected_shape)

tensorflow_graphics/rendering/utils.py

@@ -0,0 +1,40 @@
+# Copyright 2020 The TensorFlow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Various util functions common for all rasterizers."""
+
+import tensorflow as tf
+
+
+def transform_homogeneous(matrices, vertices):
+  """Applies 4x4 homogenous matrix transformations to xyz vertices.
+
+  The vertices are input and output as as row-major, but are interpreted as
+  column vectors multiplied on the right-hand side of the matrices. More
+  explicitly, this function computes (MV^T)^T where M represents transformation
+  matrices and V stands for vertices.
+  Since input vertices are xyz they are extended to xyzw with w=1.
+
+  Args:
+    matrices: A tensor of shape `[batch, 4, 4]` containing batches of view
+      projection matrices.
+    vertices: A tensor of shape `[batch, num_vertices, 3]` containing batches of
+      vertices, each defined by a 3D point.
+
+  Returns:
+    A [batch, N, 4] Tensor of xyzw vertices.
+  """
+  homogeneous_coord = tf.ones_like(vertices[..., 0:1])
+  vertices = tf.concat([vertices, homogeneous_coord], -1)
+
+  return tf.matmul(vertices, matrices, transpose_b=True)