From 6e2226baa71a2623b767358b095ae69635918490 Mon Sep 17 00:00:00 2001 From: Pavel Pidlypenskyi Date: Mon, 1 Mar 2021 16:01:49 -0800 Subject: [PATCH] Compute differentiable barycentrics. PiperOrigin-RevId: 360298407 --- tensorflow_graphics/rendering/barycentrics.py | 211 ++++++++++++++++++ .../rendering/tests/barycentrics_test.py | 55 +++++ 2 files changed, 266 insertions(+) create mode 100644 tensorflow_graphics/rendering/barycentrics.py create mode 100644 tensorflow_graphics/rendering/tests/barycentrics_test.py diff --git a/tensorflow_graphics/rendering/barycentrics.py b/tensorflow_graphics/rendering/barycentrics.py new file mode 100644 index 000000000..e5f5785cb --- /dev/null +++ b/tensorflow_graphics/rendering/barycentrics.py @@ -0,0 +1,211 @@ +# 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. +"""Set of functions to compute differentiable barycentric coordinates.""" + +from typing import Tuple + +import tensorflow as tf + +from tensorflow_graphics.rendering import framebuffer as fb +from tensorflow_graphics.util import shape + + +def differentiable_barycentrics(framebuffer: fb.Framebuffer, + clip_space_vertices: tf.Tensor, + triangles: tf.Tensor) -> fb.Framebuffer: + """Computes differentiable barycentric coordinates from a Framebuffer. + + The barycentric coordinates will be differentiable w.r.t. the input vertices. + Later, we may support derivatives w.r.t. pixel position for mip-mapping. + + Args: + framebuffer: a multi-layer Framebuffer containing triangle ids and a + foreground mask with shape [batch, num_layers, height, width, 1] + clip_space_vertices: a 2-D float32 tensor with shape [vertex_count, 4] or a + 3-D tensor with shape [batch, vertex_count, 4] containing homogenous + vertex positions (xyzw). + triangles: a 2-D int32 tensor with shape [triangle_count, 3] or a 3-D tensor + with shape [batch, triangle_count, 3] containing per-triangle vertex + indices in counter-clockwise order. + + Returns: + a copy of `framebuffer`, but the differentiable barycentric coordinates will + replace any barycentric coordinates already in the `framebuffer`. + """ + rank = lambda t: len(t.shape) + + clip_space_vertices = tf.convert_to_tensor(clip_space_vertices) + shape.check_static( + tensor=clip_space_vertices, tensor_name="clip_space_vertices", + has_rank_greater_than=1, + has_rank_less_than=4) + if rank(clip_space_vertices) == 2: + clip_space_vertices = tf.expand_dims(clip_space_vertices, axis=0) + + triangles = tf.convert_to_tensor(triangles) + shape.check_static( + tensor=triangles, tensor_name="triangles", + has_rank_greater_than=1, + has_rank_less_than=4) + if rank(triangles) == 2: + triangles = tf.expand_dims(triangles, axis=0) + + shape.compare_batch_dimensions( + tensors=(clip_space_vertices, triangles, framebuffer.triangle_id), + last_axes=(-3, -3, -4), + broadcast_compatible=False) + + # Compute image pixel coordinates. + px, py = normalized_pixel_coordinates(framebuffer.width, framebuffer.height) + + def compute_barycentrics_fn( + slices: Tuple[tf.Tensor, tf.Tensor, tf.Tensor]) -> tf.Tensor: + clip_vertices_slice, triangle_slice, triangle_id_slice = slices + triangle_id_slice = triangle_id_slice[..., 0] + if rank(triangle_id_slice) == 2: # There is no layer dimension. + triangle_id_slice = tf.expand_dims(triangle_id_slice, axis=0) + # Compute per-triangle inverse matrices. + triangle_matrices = compute_triangle_matrices(clip_vertices_slice, + triangle_slice) + + # Compute per-pixel barycentric coordinates. + barycentric_coords = compute_barycentric_coordinates( + triangle_id_slice, triangle_matrices, px, py) + barycentric_coords = tf.transpose(barycentric_coords, perm=[1, 2, 3, 0]) + return barycentric_coords + + per_image_barycentrics = tf.vectorized_map( + compute_barycentrics_fn, + (clip_space_vertices, triangles, framebuffer.triangle_id)) + + barycentric_coords = tf.stack(per_image_barycentrics, axis=0) + # After stacking barycentrics will have layers dimension no matter what. + # In order to make sure we return differentiable barycentrics of the same + # shape - reshape the tensor using original shape. + barycentric_coords = tf.reshape( + barycentric_coords, shape=framebuffer.barycentrics.value.shape) + # Mask out barycentrics for background pixels. + barycentric_coords = barycentric_coords * framebuffer.foreground_mask + + return fb.Framebuffer( + triangle_id=framebuffer.triangle_id, + vertex_ids=framebuffer.vertex_ids, + foreground_mask=framebuffer.foreground_mask, + attributes=framebuffer.attributes, + barycentrics=fb.RasterizedAttribute(barycentric_coords, None, None)) + + +def normalized_pixel_coordinates(image_width, image_height): + """Computes the normalized pixel coordinates for the specified image size. + + The x-coordinates will range from -1 to 1 left to right. + The y-coordinates will range from -1 to 1 top to bottom. + The extrema +-1 will fall onto the exterior pixel boundaries, while the + coordinates will be evaluated at pixel centers. So, image of width 4 will have + normalized pixel x-coordinates at [-0.75 -0.25 0.25 0.75], while image of + width 3 will have them at [-0.667 0 0.667]. + + Args: + image_width: int specifying desired output image width in pixels. + image_height: int specifying desired output image height in pixels. + + Returns: + Two float32 tensors with shape [image_height, image_width] containing x- and + y- coordinates, respecively, for each image pixel. + """ + width = tf.cast(image_width, tf.float32) + height = tf.cast(image_height, tf.float32) + x_range = (2 * tf.range(width) + 1) / width - 1 + y_range = (2 * tf.range(height) + 1) / height - 1 + x_coords, y_coords = tf.meshgrid(x_range, y_range) + return x_coords, y_coords + + +def compute_triangle_matrices(clip_space_vertices, triangles): + """Computes per-triangle matrices used in barycentric coordinate calculation. + + The result corresponds to the inverse matrix from equation (4) in the paper + "Triangle Scan Conversion using 2D Homogeneous Coordinates". Our matrix + inverses are not divided by the determinant, only multiplied by its sign. The + division happens in compute_barycentric_coordinates. + + Args: + clip_space_vertices: float32 tensor with shape [vertex_count, 4] containing + vertex positions in clip space (x, y, z, w). + triangles: 2-D int32 tensor with shape [triangle_count, 3]. Each triplet + contains a triangle's vertex indices in counter-clockwise order. + + Returns: + 3-D float32 tensor with shape [3, 3, triangle_count] containing per-triangle + matrices. + """ + # First make a vertex tensor of size [triangle_count, 3, 3], where the last + # dimension contains x, y, w coordinates of the corresponding vertex in each + # triangle + xyw = tf.stack([ + clip_space_vertices[:, 0], clip_space_vertices[:, 1], + clip_space_vertices[:, 3] + ], + axis=1) + xyw = tf.gather(xyw, triangles) + xyw = tf.transpose(xyw, perm=[0, 2, 1]) + # Compute the sub-determinants. + d11 = xyw[:, 1, 1] * xyw[:, 2, 2] - xyw[:, 1, 2] * xyw[:, 2, 1] + d21 = xyw[:, 1, 2] * xyw[:, 2, 0] - xyw[:, 1, 0] * xyw[:, 2, 2] + d31 = xyw[:, 1, 0] * xyw[:, 2, 1] - xyw[:, 1, 1] * xyw[:, 2, 0] + d12 = xyw[:, 2, 1] * xyw[:, 0, 2] - xyw[:, 2, 2] * xyw[:, 0, 1] + d22 = xyw[:, 2, 2] * xyw[:, 0, 0] - xyw[:, 2, 0] * xyw[:, 0, 2] + d32 = xyw[:, 2, 0] * xyw[:, 0, 1] - xyw[:, 2, 1] * xyw[:, 0, 0] + d13 = xyw[:, 0, 1] * xyw[:, 1, 2] - xyw[:, 0, 2] * xyw[:, 1, 1] + d23 = xyw[:, 0, 2] * xyw[:, 1, 0] - xyw[:, 0, 0] * xyw[:, 1, 2] + d33 = xyw[:, 0, 0] * xyw[:, 1, 1] - xyw[:, 0, 1] * xyw[:, 1, 0] + matrices = tf.stack([[d11, d12, d13], [d21, d22, d23], [d31, d32, d33]]) + # Multiply by the sign of the determinant, avoiding divide by zero. + determinant = xyw[:, 0, 0] * d11 + xyw[:, 1, 0] * d12 + xyw[:, 2, 0] * d13 + sign = tf.sign(determinant) + tf.cast(determinant == 0, tf.float32) + matrices = sign * matrices + return matrices + + +def compute_barycentric_coordinates(triangle_ids, triangle_matrices, px, py): + """Computes per-pixel barycentric coordinates. + + Args: + triangle_ids: 2-D int tensor with shape [image_height, image_width] + containing per-pixel triangle ids, as computed by rasterize_triangles. + triangle_matrices: 3-D float32 tensor with shape [3, 3, triangle_count] + containing per-triangle matrices computed by compute_triangle_matrices. + px: 2-D float32 tensor with shape [image_height, image_width] containing + per-pixel x-coordinates, as computed by normalized_pixel_coordinates. + py: 2-D float32 tensor with shape [image_height, image_width] containing + per-pixel y-coordinates, as computed by normalized_pixel_coordinates. + + Returns: + 3-D float32 tensor with shape [height, width, 3] containing the barycentric + coordinates of the point at each pixel within the triangle specified by + triangle_ids. + """ + # Gather per-pixel triangle matrices into m. + pixel_triangle_matrices = tf.gather(triangle_matrices, triangle_ids, axis=-1) + # Compute barycentric coordinates by evaluating edge equations. + barycentric_coords = ( + pixel_triangle_matrices[:, 0] * px + pixel_triangle_matrices[:, 1] * py + + pixel_triangle_matrices[:, 2]) + # Normalize so the barycentric coordinates sum to 1. Guard against division + # by zero in the case that the barycentrics sum to zero, which can happen for + # background pixels when the 0th triangle in the list is degenerate, due to + # the way we use triangle id 0 for both background and the first triangle. + barycentric_coords = tf.math.divide_no_nan( + barycentric_coords, tf.reduce_sum(barycentric_coords, axis=0)) + return barycentric_coords diff --git a/tensorflow_graphics/rendering/tests/barycentrics_test.py b/tensorflow_graphics/rendering/tests/barycentrics_test.py new file mode 100644 index 000000000..03e6d65e4 --- /dev/null +++ b/tensorflow_graphics/rendering/tests/barycentrics_test.py @@ -0,0 +1,55 @@ +# 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 differentiable barycentrics.""" + +from absl.testing import parameterized +import tensorflow as tf + +from tensorflow_graphics.rendering import barycentrics +from tensorflow_graphics.rendering import rasterization_backend +from tensorflow_graphics.rendering import utils +from tensorflow_graphics.util import test_case + + +class BarycentricsTest(test_case.TestCase): + + @parameterized.parameters( + (rasterization_backend.RasterizationBackends.CPU,)) + def test_differentialbe_barycentrics_close_to_rasterizer( + self, rasterization_backend_type): + image_width = 32 + image_height = 32 + + initial_vertices = tf.constant([[[0, 0, 0], [0.5, 0, 0], [0.5, 0.5, 0]]], + dtype=tf.float32) + triangles = [[0, 1, 2]] + view_projection_matrix = tf.expand_dims(tf.eye(4), axis=0) + + # Compute rasterization barycentrics + rasterized = rasterization_backend.rasterize( + initial_vertices, + triangles, + view_projection_matrix, (image_width, image_height), + num_layers=1, + backend=rasterization_backend_type) + + ras_barycentrics = rasterized.barycentrics.value + + clip_space_vertices = utils.transform_homogeneous(view_projection_matrix, + initial_vertices) + rasterized_w_diff_barycentrics = barycentrics.differentiable_barycentrics( + rasterized, clip_space_vertices, triangles) + diff_barycentrics = rasterized_w_diff_barycentrics.barycentrics.value + + self.assertAllClose(ras_barycentrics, diff_barycentrics, rtol=1e-4)