We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Dense Image Warp tests are flaky
Recently we've seen that #53 is causing flaky failures in the CI. See:
https://source.cloud.google.com/results/invocations/8f31faef-505a-440e-b75f-e6edf1071269/targets/tensorflow_addons%2Fubuntu%2Fgpu%2Fpy3%2Fpresubmit/log

Do you mind taking a look when time allows @WindQAQ ?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/image/dense_image_warp.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 """Image warping using per-pixel flow vectors."""
16 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 
20 import numpy as np
21 import tensorflow as tf
22 
23 
24 @tf.function
25 def interpolate_bilinear(grid, query_points, indexing="ij", name=None):
26     """Similar to Matlab's interp2 function.
27 
28     Finds values for query points on a grid using bilinear interpolation.
29 
30     Args:
31       grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.
32       query_points: a 3-D float `Tensor` of N points with shape
33         `[batch, N, 2]`.
34       indexing: whether the query points are specified as row and column (ij),
35         or Cartesian coordinates (xy).
36       name: a name for the operation (optional).
37 
38     Returns:
39       values: a 3-D `Tensor` with shape `[batch, N, channels]`
40 
41     Raises:
42       ValueError: if the indexing mode is invalid, or if the shape of the
43         inputs invalid.
44     """
45     if indexing != "ij" and indexing != "xy":
46         raise ValueError("Indexing mode must be \'ij\' or \'xy\'")
47 
48     with tf.name_scope(name or "interpolate_bilinear"):
49         grid = tf.convert_to_tensor(grid)
50         query_points = tf.convert_to_tensor(query_points)
51         shape = grid.get_shape().as_list()
52         if len(shape) != 4:
53             msg = "Grid must be 4 dimensional. Received size: "
54             raise ValueError(msg + str(grid.get_shape()))
55 
56         batch_size, height, width, channels = (tf.shape(grid)[0],
57                                                tf.shape(grid)[1],
58                                                tf.shape(grid)[2],
59                                                tf.shape(grid)[3])
60 
61         shape = [batch_size, height, width, channels]
62         query_type = query_points.dtype
63         grid_type = grid.dtype
64 
65         tf.debugging.assert_equal(
66             len(query_points.get_shape()),
67             3,
68             message="Query points must be 3 dimensional.")
69         tf.debugging.assert_equal(
70             tf.shape(query_points)[2],
71             2,
72             message="Query points must be size 2 in dim 2.")
73 
74         num_queries = tf.shape(query_points)[1]
75 
76         tf.debugging.assert_greater_equal(
77             height, 2, message="Grid height must be at least 2."),
78         tf.debugging.assert_greater_equal(
79             width, 2, message="Grid width must be at least 2.")
80 
81         alphas = []
82         floors = []
83         ceils = []
84         index_order = [0, 1] if indexing == "ij" else [1, 0]
85         unstacked_query_points = tf.unstack(query_points, axis=2)
86 
87         for dim in index_order:
88             with tf.name_scope("dim-" + str(dim)):
89                 queries = unstacked_query_points[dim]
90 
91                 size_in_indexing_dimension = shape[dim + 1]
92 
93                 # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1
94                 # is still a valid index into the grid.
95                 max_floor = tf.cast(size_in_indexing_dimension - 2, query_type)
96                 min_floor = tf.constant(0.0, dtype=query_type)
97                 floor = tf.math.minimum(
98                     tf.math.maximum(min_floor, tf.math.floor(queries)),
99                     max_floor)
100                 int_floor = tf.cast(floor, tf.dtypes.int32)
101                 floors.append(int_floor)
102                 ceil = int_floor + 1
103                 ceils.append(ceil)
104 
105                 # alpha has the same type as the grid, as we will directly use alpha
106                 # when taking linear combinations of pixel values from the image.
107                 alpha = tf.cast(queries - floor, grid_type)
108                 min_alpha = tf.constant(0.0, dtype=grid_type)
109                 max_alpha = tf.constant(1.0, dtype=grid_type)
110                 alpha = tf.math.minimum(
111                     tf.math.maximum(min_alpha, alpha), max_alpha)
112 
113                 # Expand alpha to [b, n, 1] so we can use broadcasting
114                 # (since the alpha values don't depend on the channel).
115                 alpha = tf.expand_dims(alpha, 2)
116                 alphas.append(alpha)
117 
118         tf.debugging.assert_less_equal(
119             tf.cast(batch_size * height * width, dtype=tf.dtypes.float32),
120             np.iinfo(np.int32).max / 8.0,
121             message="The image size or batch size is sufficiently large "
122             "that the linearized addresses used by tf.gather "
123             "may exceed the int32 limit.")
124         flattened_grid = tf.reshape(grid,
125                                     [batch_size * height * width, channels])
126         batch_offsets = tf.reshape(
127             tf.range(batch_size) * height * width, [batch_size, 1])
128 
129         # This wraps tf.gather. We reshape the image data such that the
130         # batch, y, and x coordinates are pulled into the first dimension.
131         # Then we gather. Finally, we reshape the output back. It's possible this
132         # code would be made simpler by using tf.gather_nd.
133         def gather(y_coords, x_coords, name):
134             with tf.name_scope("gather-" + name):
135                 linear_coordinates = (
136                     batch_offsets + y_coords * width + x_coords)
137                 gathered_values = tf.gather(flattened_grid, linear_coordinates)
138                 return tf.reshape(gathered_values,
139                                   [batch_size, num_queries, channels])
140 
141         # grab the pixel values in the 4 corners around each query point
142         top_left = gather(floors[0], floors[1], "top_left")
143         top_right = gather(floors[0], ceils[1], "top_right")
144         bottom_left = gather(ceils[0], floors[1], "bottom_left")
145         bottom_right = gather(ceils[0], ceils[1], "bottom_right")
146 
147         # now, do the actual interpolation
148         with tf.name_scope("interpolate"):
149             interp_top = alphas[1] * (top_right - top_left) + top_left
150             interp_bottom = alphas[1] * (
151                 bottom_right - bottom_left) + bottom_left
152             interp = alphas[0] * (interp_bottom - interp_top) + interp_top
153 
154         return interp
155 
156 
157 @tf.function
158 def dense_image_warp(image, flow, name=None):
159     """Image warping using per-pixel flow vectors.
160 
161     Apply a non-linear warp to the image, where the warp is specified by a
162     dense flow field of offset vectors that define the correspondences of
163     pixel values in the output image back to locations in the source image.
164     Specifically, the pixel value at output[b, j, i, c] is
165     images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].
166 
167     The locations specified by this formula do not necessarily map to an int
168     index. Therefore, the pixel value is obtained by bilinear
169     interpolation of the 4 nearest pixels around
170     (b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside
171     of the image, we use the nearest pixel values at the image boundary.
172 
173     Args:
174       image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
175       flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
176       name: A name for the operation (optional).
177 
178       Note that image and flow can be of type tf.half, tf.float32, or
179       tf.float64, and do not necessarily have to be the same type.
180 
181     Returns:
182       A 4-D float `Tensor` with shape`[batch, height, width, channels]`
183         and same type as input image.
184 
185     Raises:
186       ValueError: if height < 2 or width < 2 or the inputs have the wrong
187         number of dimensions.
188     """
189     with tf.name_scope(name or "dense_image_warp"):
190         image = tf.convert_to_tensor(image)
191         flow = tf.convert_to_tensor(flow)
192         batch_size, height, width, channels = (tf.shape(image)[0],
193                                                tf.shape(image)[1],
194                                                tf.shape(image)[2],
195                                                tf.shape(image)[3])
196 
197         # The flow is defined on the image grid. Turn the flow into a list of query
198         # points in the grid space.
199         grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))
200         stacked_grid = tf.cast(tf.stack([grid_y, grid_x], axis=2), flow.dtype)
201         batched_grid = tf.expand_dims(stacked_grid, axis=0)
202         query_points_on_grid = batched_grid - flow
203         query_points_flattened = tf.reshape(query_points_on_grid,
204                                             [batch_size, height * width, 2])
205         # Compute values at the query points, then reshape the result back to the
206         # image grid.
207         interpolated = interpolate_bilinear(image, query_points_flattened)
208         interpolated = tf.reshape(interpolated,
209                                   [batch_size, height, width, channels])
210         return interpolated
211
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/tensorflow_addons/image/dense_image_warp.py b/tensorflow_addons/image/dense_image_warp.py
--- a/tensorflow_addons/image/dense_image_warp.py
+++ b/tensorflow_addons/image/dense_image_warp.py
@@ -21,7 +21,6 @@
 import tensorflow as tf
 
 
-@tf.function
 def interpolate_bilinear(grid, query_points, indexing="ij", name=None):
     """Similar to Matlab's interp2 function.
 
@@ -48,30 +47,28 @@
     with tf.name_scope(name or "interpolate_bilinear"):
         grid = tf.convert_to_tensor(grid)
         query_points = tf.convert_to_tensor(query_points)
-        shape = grid.get_shape().as_list()
-        if len(shape) != 4:
+
+        if len(grid.shape) != 4:
             msg = "Grid must be 4 dimensional. Received size: "
-            raise ValueError(msg + str(grid.get_shape()))
+            raise ValueError(msg + str(grid.shape))
+
+        if len(query_points.shape) != 3:
+            raise ValueError("Query points must be 3 dimensional.")
+
+        grid_shape = tf.shape(grid)
+        query_shape = tf.shape(query_points)
 
-        batch_size, height, width, channels = (tf.shape(grid)[0],
-                                               tf.shape(grid)[1],
-                                               tf.shape(grid)[2],
-                                               tf.shape(grid)[3])
+        batch_size, height, width, channels = (grid_shape[0], grid_shape[1],
+                                               grid_shape[2], grid_shape[3])
 
         shape = [batch_size, height, width, channels]
+        num_queries = query_shape[1]
+
         query_type = query_points.dtype
         grid_type = grid.dtype
 
         tf.debugging.assert_equal(
-            len(query_points.get_shape()),
-            3,
-            message="Query points must be 3 dimensional.")
-        tf.debugging.assert_equal(
-            tf.shape(query_points)[2],
-            2,
-            message="Query points must be size 2 in dim 2.")
-
-        num_queries = tf.shape(query_points)[1]
+            query_shape[2], 2, message="Query points must be size 2 in dim 2.")
 
         tf.debugging.assert_greater_equal(
             height, 2, message="Grid height must be at least 2."),

{"golden_diff": "diff --git a/tensorflow_addons/image/dense_image_warp.py b/tensorflow_addons/image/dense_image_warp.py\n--- a/tensorflow_addons/image/dense_image_warp.py\n+++ b/tensorflow_addons/image/dense_image_warp.py\n@@ -21,7 +21,6 @@\n import tensorflow as tf\n \n \n-@tf.function\n def interpolate_bilinear(grid, query_points, indexing=\"ij\", name=None):\n     \"\"\"Similar to Matlab's interp2 function.\n \n@@ -48,30 +47,28 @@\n     with tf.name_scope(name or \"interpolate_bilinear\"):\n         grid = tf.convert_to_tensor(grid)\n         query_points = tf.convert_to_tensor(query_points)\n-        shape = grid.get_shape().as_list()\n-        if len(shape) != 4:\n+\n+        if len(grid.shape) != 4:\n             msg = \"Grid must be 4 dimensional. Received size: \"\n-            raise ValueError(msg + str(grid.get_shape()))\n+            raise ValueError(msg + str(grid.shape))\n+\n+        if len(query_points.shape) != 3:\n+            raise ValueError(\"Query points must be 3 dimensional.\")\n+\n+        grid_shape = tf.shape(grid)\n+        query_shape = tf.shape(query_points)\n \n-        batch_size, height, width, channels = (tf.shape(grid)[0],\n-                                               tf.shape(grid)[1],\n-                                               tf.shape(grid)[2],\n-                                               tf.shape(grid)[3])\n+        batch_size, height, width, channels = (grid_shape[0], grid_shape[1],\n+                                               grid_shape[2], grid_shape[3])\n \n         shape = [batch_size, height, width, channels]\n+        num_queries = query_shape[1]\n+\n         query_type = query_points.dtype\n         grid_type = grid.dtype\n \n         tf.debugging.assert_equal(\n-            len(query_points.get_shape()),\n-            3,\n-            message=\"Query points must be 3 dimensional.\")\n-        tf.debugging.assert_equal(\n-            tf.shape(query_points)[2],\n-            2,\n-            message=\"Query points must be size 2 in dim 2.\")\n-\n-        num_queries = tf.shape(query_points)[1]\n+            query_shape[2], 2, message=\"Query points must be size 2 in dim 2.\")\n \n         tf.debugging.assert_greater_equal(\n             height, 2, message=\"Grid height must be at least 2.\"),\n", "issue": "Dense Image Warp tests are flaky\nRecently we've seen that #53 is causing flaky failures in the CI. See:\r\nhttps://source.cloud.google.com/results/invocations/8f31faef-505a-440e-b75f-e6edf1071269/targets/tensorflow_addons%2Fubuntu%2Fgpu%2Fpy3%2Fpresubmit/log\r\n\r\nDo you mind taking a look when time allows @WindQAQ ?\n", "before_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Image warping using per-pixel flow vectors.\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\nimport tensorflow as tf\n\n\n@tf.function\ndef interpolate_bilinear(grid, query_points, indexing=\"ij\", name=None):\n    \"\"\"Similar to Matlab's interp2 function.\n\n    Finds values for query points on a grid using bilinear interpolation.\n\n    Args:\n      grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.\n      query_points: a 3-D float `Tensor` of N points with shape\n        `[batch, N, 2]`.\n      indexing: whether the query points are specified as row and column (ij),\n        or Cartesian coordinates (xy).\n      name: a name for the operation (optional).\n\n    Returns:\n      values: a 3-D `Tensor` with shape `[batch, N, channels]`\n\n    Raises:\n      ValueError: if the indexing mode is invalid, or if the shape of the\n        inputs invalid.\n    \"\"\"\n    if indexing != \"ij\" and indexing != \"xy\":\n        raise ValueError(\"Indexing mode must be \\'ij\\' or \\'xy\\'\")\n\n    with tf.name_scope(name or \"interpolate_bilinear\"):\n        grid = tf.convert_to_tensor(grid)\n        query_points = tf.convert_to_tensor(query_points)\n        shape = grid.get_shape().as_list()\n        if len(shape) != 4:\n            msg = \"Grid must be 4 dimensional. Received size: \"\n            raise ValueError(msg + str(grid.get_shape()))\n\n        batch_size, height, width, channels = (tf.shape(grid)[0],\n                                               tf.shape(grid)[1],\n                                               tf.shape(grid)[2],\n                                               tf.shape(grid)[3])\n\n        shape = [batch_size, height, width, channels]\n        query_type = query_points.dtype\n        grid_type = grid.dtype\n\n        tf.debugging.assert_equal(\n            len(query_points.get_shape()),\n            3,\n            message=\"Query points must be 3 dimensional.\")\n        tf.debugging.assert_equal(\n            tf.shape(query_points)[2],\n            2,\n            message=\"Query points must be size 2 in dim 2.\")\n\n        num_queries = tf.shape(query_points)[1]\n\n        tf.debugging.assert_greater_equal(\n            height, 2, message=\"Grid height must be at least 2.\"),\n        tf.debugging.assert_greater_equal(\n            width, 2, message=\"Grid width must be at least 2.\")\n\n        alphas = []\n        floors = []\n        ceils = []\n        index_order = [0, 1] if indexing == \"ij\" else [1, 0]\n        unstacked_query_points = tf.unstack(query_points, axis=2)\n\n        for dim in index_order:\n            with tf.name_scope(\"dim-\" + str(dim)):\n                queries = unstacked_query_points[dim]\n\n                size_in_indexing_dimension = shape[dim + 1]\n\n                # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1\n                # is still a valid index into the grid.\n                max_floor = tf.cast(size_in_indexing_dimension - 2, query_type)\n                min_floor = tf.constant(0.0, dtype=query_type)\n                floor = tf.math.minimum(\n                    tf.math.maximum(min_floor, tf.math.floor(queries)),\n                    max_floor)\n                int_floor = tf.cast(floor, tf.dtypes.int32)\n                floors.append(int_floor)\n                ceil = int_floor + 1\n                ceils.append(ceil)\n\n                # alpha has the same type as the grid, as we will directly use alpha\n                # when taking linear combinations of pixel values from the image.\n                alpha = tf.cast(queries - floor, grid_type)\n                min_alpha = tf.constant(0.0, dtype=grid_type)\n                max_alpha = tf.constant(1.0, dtype=grid_type)\n                alpha = tf.math.minimum(\n                    tf.math.maximum(min_alpha, alpha), max_alpha)\n\n                # Expand alpha to [b, n, 1] so we can use broadcasting\n                # (since the alpha values don't depend on the channel).\n                alpha = tf.expand_dims(alpha, 2)\n                alphas.append(alpha)\n\n        tf.debugging.assert_less_equal(\n            tf.cast(batch_size * height * width, dtype=tf.dtypes.float32),\n            np.iinfo(np.int32).max / 8.0,\n            message=\"The image size or batch size is sufficiently large \"\n            \"that the linearized addresses used by tf.gather \"\n            \"may exceed the int32 limit.\")\n        flattened_grid = tf.reshape(grid,\n                                    [batch_size * height * width, channels])\n        batch_offsets = tf.reshape(\n            tf.range(batch_size) * height * width, [batch_size, 1])\n\n        # This wraps tf.gather. We reshape the image data such that the\n        # batch, y, and x coordinates are pulled into the first dimension.\n        # Then we gather. Finally, we reshape the output back. It's possible this\n        # code would be made simpler by using tf.gather_nd.\n        def gather(y_coords, x_coords, name):\n            with tf.name_scope(\"gather-\" + name):\n                linear_coordinates = (\n                    batch_offsets + y_coords * width + x_coords)\n                gathered_values = tf.gather(flattened_grid, linear_coordinates)\n                return tf.reshape(gathered_values,\n                                  [batch_size, num_queries, channels])\n\n        # grab the pixel values in the 4 corners around each query point\n        top_left = gather(floors[0], floors[1], \"top_left\")\n        top_right = gather(floors[0], ceils[1], \"top_right\")\n        bottom_left = gather(ceils[0], floors[1], \"bottom_left\")\n        bottom_right = gather(ceils[0], ceils[1], \"bottom_right\")\n\n        # now, do the actual interpolation\n        with tf.name_scope(\"interpolate\"):\n            interp_top = alphas[1] * (top_right - top_left) + top_left\n            interp_bottom = alphas[1] * (\n                bottom_right - bottom_left) + bottom_left\n            interp = alphas[0] * (interp_bottom - interp_top) + interp_top\n\n        return interp\n\n\n@tf.function\ndef dense_image_warp(image, flow, name=None):\n    \"\"\"Image warping using per-pixel flow vectors.\n\n    Apply a non-linear warp to the image, where the warp is specified by a\n    dense flow field of offset vectors that define the correspondences of\n    pixel values in the output image back to locations in the source image.\n    Specifically, the pixel value at output[b, j, i, c] is\n    images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].\n\n    The locations specified by this formula do not necessarily map to an int\n    index. Therefore, the pixel value is obtained by bilinear\n    interpolation of the 4 nearest pixels around\n    (b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside\n    of the image, we use the nearest pixel values at the image boundary.\n\n    Args:\n      image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.\n      flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.\n      name: A name for the operation (optional).\n\n      Note that image and flow can be of type tf.half, tf.float32, or\n      tf.float64, and do not necessarily have to be the same type.\n\n    Returns:\n      A 4-D float `Tensor` with shape`[batch, height, width, channels]`\n        and same type as input image.\n\n    Raises:\n      ValueError: if height < 2 or width < 2 or the inputs have the wrong\n        number of dimensions.\n    \"\"\"\n    with tf.name_scope(name or \"dense_image_warp\"):\n        image = tf.convert_to_tensor(image)\n        flow = tf.convert_to_tensor(flow)\n        batch_size, height, width, channels = (tf.shape(image)[0],\n                                               tf.shape(image)[1],\n                                               tf.shape(image)[2],\n                                               tf.shape(image)[3])\n\n        # The flow is defined on the image grid. Turn the flow into a list of query\n        # points in the grid space.\n        grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))\n        stacked_grid = tf.cast(tf.stack([grid_y, grid_x], axis=2), flow.dtype)\n        batched_grid = tf.expand_dims(stacked_grid, axis=0)\n        query_points_on_grid = batched_grid - flow\n        query_points_flattened = tf.reshape(query_points_on_grid,\n                                            [batch_size, height * width, 2])\n        # Compute values at the query points, then reshape the result back to the\n        # image grid.\n        interpolated = interpolate_bilinear(image, query_points_flattened)\n        interpolated = tf.reshape(interpolated,\n                                  [batch_size, height, width, channels])\n        return interpolated\n", "path": "tensorflow_addons/image/dense_image_warp.py"}], "after_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Image warping using per-pixel flow vectors.\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\nimport tensorflow as tf\n\n\ndef interpolate_bilinear(grid, query_points, indexing=\"ij\", name=None):\n    \"\"\"Similar to Matlab's interp2 function.\n\n    Finds values for query points on a grid using bilinear interpolation.\n\n    Args:\n      grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.\n      query_points: a 3-D float `Tensor` of N points with shape\n        `[batch, N, 2]`.\n      indexing: whether the query points are specified as row and column (ij),\n        or Cartesian coordinates (xy).\n      name: a name for the operation (optional).\n\n    Returns:\n      values: a 3-D `Tensor` with shape `[batch, N, channels]`\n\n    Raises:\n      ValueError: if the indexing mode is invalid, or if the shape of the\n        inputs invalid.\n    \"\"\"\n    if indexing != \"ij\" and indexing != \"xy\":\n        raise ValueError(\"Indexing mode must be \\'ij\\' or \\'xy\\'\")\n\n    with tf.name_scope(name or \"interpolate_bilinear\"):\n        grid = tf.convert_to_tensor(grid)\n        query_points = tf.convert_to_tensor(query_points)\n\n        if len(grid.shape) != 4:\n            msg = \"Grid must be 4 dimensional. Received size: \"\n            raise ValueError(msg + str(grid.shape))\n\n        if len(query_points.shape) != 3:\n            raise ValueError(\"Query points must be 3 dimensional.\")\n\n        grid_shape = tf.shape(grid)\n        query_shape = tf.shape(query_points)\n\n        batch_size, height, width, channels = (grid_shape[0], grid_shape[1],\n                                               grid_shape[2], grid_shape[3])\n\n        shape = [batch_size, height, width, channels]\n        num_queries = query_shape[1]\n\n        query_type = query_points.dtype\n        grid_type = grid.dtype\n\n        tf.debugging.assert_equal(\n            query_shape[2], 2, message=\"Query points must be size 2 in dim 2.\")\n\n        tf.debugging.assert_greater_equal(\n            height, 2, message=\"Grid height must be at least 2.\"),\n        tf.debugging.assert_greater_equal(\n            width, 2, message=\"Grid width must be at least 2.\")\n\n        alphas = []\n        floors = []\n        ceils = []\n        index_order = [0, 1] if indexing == \"ij\" else [1, 0]\n        unstacked_query_points = tf.unstack(query_points, axis=2)\n\n        for dim in index_order:\n            with tf.name_scope(\"dim-\" + str(dim)):\n                queries = unstacked_query_points[dim]\n\n                size_in_indexing_dimension = shape[dim + 1]\n\n                # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1\n                # is still a valid index into the grid.\n                max_floor = tf.cast(size_in_indexing_dimension - 2, query_type)\n                min_floor = tf.constant(0.0, dtype=query_type)\n                floor = tf.math.minimum(\n                    tf.math.maximum(min_floor, tf.math.floor(queries)),\n                    max_floor)\n                int_floor = tf.cast(floor, tf.dtypes.int32)\n                floors.append(int_floor)\n                ceil = int_floor + 1\n                ceils.append(ceil)\n\n                # alpha has the same type as the grid, as we will directly use alpha\n                # when taking linear combinations of pixel values from the image.\n                alpha = tf.cast(queries - floor, grid_type)\n                min_alpha = tf.constant(0.0, dtype=grid_type)\n                max_alpha = tf.constant(1.0, dtype=grid_type)\n                alpha = tf.math.minimum(\n                    tf.math.maximum(min_alpha, alpha), max_alpha)\n\n                # Expand alpha to [b, n, 1] so we can use broadcasting\n                # (since the alpha values don't depend on the channel).\n                alpha = tf.expand_dims(alpha, 2)\n                alphas.append(alpha)\n\n        tf.debugging.assert_less_equal(\n            tf.cast(batch_size * height * width, dtype=tf.dtypes.float32),\n            np.iinfo(np.int32).max / 8.0,\n            message=\"The image size or batch size is sufficiently large \"\n            \"that the linearized addresses used by tf.gather \"\n            \"may exceed the int32 limit.\")\n        flattened_grid = tf.reshape(grid,\n                                    [batch_size * height * width, channels])\n        batch_offsets = tf.reshape(\n            tf.range(batch_size) * height * width, [batch_size, 1])\n\n        # This wraps tf.gather. We reshape the image data such that the\n        # batch, y, and x coordinates are pulled into the first dimension.\n        # Then we gather. Finally, we reshape the output back. It's possible this\n        # code would be made simpler by using tf.gather_nd.\n        def gather(y_coords, x_coords, name):\n            with tf.name_scope(\"gather-\" + name):\n                linear_coordinates = (\n                    batch_offsets + y_coords * width + x_coords)\n                gathered_values = tf.gather(flattened_grid, linear_coordinates)\n                return tf.reshape(gathered_values,\n                                  [batch_size, num_queries, channels])\n\n        # grab the pixel values in the 4 corners around each query point\n        top_left = gather(floors[0], floors[1], \"top_left\")\n        top_right = gather(floors[0], ceils[1], \"top_right\")\n        bottom_left = gather(ceils[0], floors[1], \"bottom_left\")\n        bottom_right = gather(ceils[0], ceils[1], \"bottom_right\")\n\n        # now, do the actual interpolation\n        with tf.name_scope(\"interpolate\"):\n            interp_top = alphas[1] * (top_right - top_left) + top_left\n            interp_bottom = alphas[1] * (\n                bottom_right - bottom_left) + bottom_left\n            interp = alphas[0] * (interp_bottom - interp_top) + interp_top\n\n        return interp\n\n\n@tf.function\ndef dense_image_warp(image, flow, name=None):\n    \"\"\"Image warping using per-pixel flow vectors.\n\n    Apply a non-linear warp to the image, where the warp is specified by a\n    dense flow field of offset vectors that define the correspondences of\n    pixel values in the output image back to locations in the source image.\n    Specifically, the pixel value at output[b, j, i, c] is\n    images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].\n\n    The locations specified by this formula do not necessarily map to an int\n    index. Therefore, the pixel value is obtained by bilinear\n    interpolation of the 4 nearest pixels around\n    (b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside\n    of the image, we use the nearest pixel values at the image boundary.\n\n    Args:\n      image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.\n      flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.\n      name: A name for the operation (optional).\n\n      Note that image and flow can be of type tf.half, tf.float32, or\n      tf.float64, and do not necessarily have to be the same type.\n\n    Returns:\n      A 4-D float `Tensor` with shape`[batch, height, width, channels]`\n        and same type as input image.\n\n    Raises:\n      ValueError: if height < 2 or width < 2 or the inputs have the wrong\n        number of dimensions.\n    \"\"\"\n    with tf.name_scope(name or \"dense_image_warp\"):\n        image = tf.convert_to_tensor(image)\n        flow = tf.convert_to_tensor(flow)\n        batch_size, height, width, channels = (tf.shape(image)[0],\n                                               tf.shape(image)[1],\n                                               tf.shape(image)[2],\n                                               tf.shape(image)[3])\n\n        # The flow is defined on the image grid. Turn the flow into a list of query\n        # points in the grid space.\n        grid_x, grid_y = tf.meshgrid(tf.range(width), tf.range(height))\n        stacked_grid = tf.cast(tf.stack([grid_y, grid_x], axis=2), flow.dtype)\n        batched_grid = tf.expand_dims(stacked_grid, axis=0)\n        query_points_on_grid = batched_grid - flow\n        query_points_flattened = tf.reshape(query_points_on_grid,\n                                            [batch_size, height * width, 2])\n        # Compute values at the query points, then reshape the result back to the\n        # image grid.\n        interpolated = interpolate_bilinear(image, query_points_flattened)\n        interpolated = tf.reshape(interpolated,\n                                  [batch_size, height, width, channels])\n        return interpolated\n", "path": "tensorflow_addons/image/dense_image_warp.py"}]}