Add unit test verifying tf.Dataset support for dense_image_warp (#332) (#654)

Author: mels630

tensorflow_addons/image/dense_image_warp_test.py

@@ -154,16 +154,25 @@ def _check_interpolation_correctness(self,
                                          shape,
                                          image_type,
                                          flow_type,
+                                         call_with_unknown_shapes=False,
                                          num_probes=5):
         """Interpolate, and then assert correctness for a few query
         locations."""
         low_precision = image_type == "float16" or flow_type == "float16"
         rand_image, rand_flows = self._get_random_image_and_flows(
             shape, image_type, flow_type)
 
-        interp = dense_image_warp(
-            image=tf.convert_to_tensor(rand_image),
-            flow=tf.convert_to_tensor(rand_flows))
+        if call_with_unknown_shapes:
+            fn = dense_image_warp.get_concrete_function(
+                tf.TensorSpec(shape=None, dtype=image_type),
+                tf.TensorSpec(shape=None, dtype=flow_type))
+            interp = fn(
+                image=tf.convert_to_tensor(rand_image),
+                flow=tf.convert_to_tensor(rand_flows))
+        else:
+            interp = dense_image_warp(
+                image=tf.convert_to_tensor(rand_image),
+                flow=tf.convert_to_tensor(rand_flows))
 
         for _ in range(num_probes):
             batch_index = np.random.randint(0, shape[0])
@@ -189,6 +198,14 @@ def test_interpolation(self):
                     self._check_interpolation_correctness(
                         shape, im_type, flow_type)
 
+    def test_unknown_shapes(self):
+        """Apply _check_interpolation_correctness() for a few sizes and check
+        for tf.Dataset compatibility."""
+        shapes_to_try = [[3, 4, 5, 6], [1, 5, 5, 3], [1, 2, 2, 1]]
+        for shape in shapes_to_try:
+            self._check_interpolation_correctness(shape, "float32", "float32",
+                                                  True)
+
     def test_gradients_exist(self):
         """Check that backprop can run.