Used fixture instead. (#1331)

Author: gabrieldemarmiesse

tensorflow_addons/image/dense_image_warp_test.py

@@ -90,66 +90,9 @@ def test_unknown_shape(self):
 
 @test_utils.run_all_in_graph_and_eager_modes
 class DenseImageWarpTest(tf.test.TestCase):
-    def _get_random_image_and_flows(self, shape, image_type, flow_type):
-        batch_size, height, width, num_channels = shape
-        image_shape = [batch_size, height, width, num_channels]
-        image = np.random.normal(size=image_shape)
-        flow_shape = [batch_size, height, width, 2]
-        flows = np.random.normal(size=flow_shape) * 3
-        return image.astype(image_type), flows.astype(flow_type)
-
-    def _assert_correct_interpolation_value(
-        self,
-        image,
-        flows,
-        pred_interpolation,
-        batch_index,
-        y_index,
-        x_index,
-        low_precision=False,
-    ):
-        """Assert that the tf interpolation matches hand-computed value."""
-        height = image.shape[1]
-        width = image.shape[2]
-        displacement = flows[batch_index, y_index, x_index, :]
-        float_y = y_index - displacement[0]
-        float_x = x_index - displacement[1]
-        floor_y = max(min(height - 2, math.floor(float_y)), 0)
-        floor_x = max(min(width - 2, math.floor(float_x)), 0)
-        ceil_y = floor_y + 1
-        ceil_x = floor_x + 1
-
-        alpha_y = min(max(0.0, float_y - floor_y), 1.0)
-        alpha_x = min(max(0.0, float_x - floor_x), 1.0)
-
-        floor_y = int(floor_y)
-        floor_x = int(floor_x)
-        ceil_y = int(ceil_y)
-        ceil_x = int(ceil_x)
-
-        top_left = image[batch_index, floor_y, floor_x, :]
-        top_right = image[batch_index, floor_y, ceil_x, :]
-        bottom_left = image[batch_index, ceil_y, floor_x, :]
-        bottom_right = image[batch_index, ceil_y, ceil_x, :]
-
-        interp_top = alpha_x * (top_right - top_left) + top_left
-        interp_bottom = alpha_x * (bottom_right - bottom_left) + bottom_left
-        interp = alpha_y * (interp_bottom - interp_top) + interp_top
-        atol = 1e-6
-        rtol = 1e-6
-        if low_precision:
-            atol = 1e-2
-            rtol = 1e-3
-        self.assertAllClose(
-            interp,
-            pred_interpolation[batch_index, y_index, x_index, :],
-            atol=atol,
-            rtol=rtol,
-        )
-
     def _check_zero_flow_correctness(self, shape, image_type, flow_type):
         """Assert using zero flows doesn't change the input image."""
-        rand_image, rand_flows = self._get_random_image_and_flows(
+        rand_image, rand_flows = _get_random_image_and_flows(
             shape, image_type, flow_type
         )
         rand_flows *= 0
@@ -169,62 +112,6 @@ def test_zero_flows(self):
                 shape, image_type="float32", flow_type="float32"
             )
 
-    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
-        )
-
-        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])
-            y_index = np.random.randint(0, shape[1])
-            x_index = np.random.randint(0, shape[2])
-
-            self._assert_correct_interpolation_value(
-                rand_image,
-                rand_flows,
-                interp,
-                batch_index,
-                y_index,
-                x_index,
-                low_precision=low_precision,
-            )
-
-    def test_interpolation(self):
-        """Apply _check_interpolation_correctness() for a few sizes and
-        types."""
-        shapes_to_try = [[3, 4, 5, 6], [1, 2, 2, 1]]
-        for im_type in ["float32", "float64", "float16"]:
-            for flow_type in ["float32", "float64", "float16"]:
-                for shape in shapes_to_try:
-                    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, 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.
 
@@ -253,12 +140,132 @@ def loss():
         for _ in range(10):
             self.evaluate(minimize_op)
 
-    def test_size_exception(self):
-        """Make sure it throws an exception for images that are too small."""
-        shape = [1, 2, 1, 1]
-        errors = (ValueError, tf.errors.InvalidArgumentError)
-        with self.assertRaisesRegexp(errors, "Grid width must be at least 2."):
-            self._check_interpolation_correctness(shape, "float32", "float32")
+
+def _assert_correct_interpolation_value(
+    image,
+    flows,
+    pred_interpolation,
+    batch_index,
+    y_index,
+    x_index,
+    low_precision=False,
+):
+    """Assert that the tf interpolation matches hand-computed value."""
+    height = image.shape[1]
+    width = image.shape[2]
+    displacement = flows[batch_index, y_index, x_index, :]
+    float_y = y_index - displacement[0]
+    float_x = x_index - displacement[1]
+    floor_y = max(min(height - 2, math.floor(float_y)), 0)
+    floor_x = max(min(width - 2, math.floor(float_x)), 0)
+    ceil_y = floor_y + 1
+    ceil_x = floor_x + 1
+
+    alpha_y = min(max(0.0, float_y - floor_y), 1.0)
+    alpha_x = min(max(0.0, float_x - floor_x), 1.0)
+
+    floor_y = int(floor_y)
+    floor_x = int(floor_x)
+    ceil_y = int(ceil_y)
+    ceil_x = int(ceil_x)
+
+    top_left = image[batch_index, floor_y, floor_x, :]
+    top_right = image[batch_index, floor_y, ceil_x, :]
+    bottom_left = image[batch_index, ceil_y, floor_x, :]
+    bottom_right = image[batch_index, ceil_y, ceil_x, :]
+
+    interp_top = alpha_x * (top_right - top_left) + top_left
+    interp_bottom = alpha_x * (bottom_right - bottom_left) + bottom_left
+    interp = alpha_y * (interp_bottom - interp_top) + interp_top
+    atol = 1e-6
+    rtol = 1e-6
+    if low_precision:
+        atol = 1e-2
+        rtol = 1e-3
+    np.testing.assert_allclose(
+        interp,
+        pred_interpolation[batch_index, y_index, x_index, :],
+        atol=atol,
+        rtol=rtol,
+    )
+
+
+def _get_random_image_and_flows(shape, image_type, flow_type):
+    batch_size, height, width, num_channels = shape
+    image_shape = [batch_size, height, width, num_channels]
+    image = np.random.normal(size=image_shape)
+    flow_shape = [batch_size, height, width, 2]
+    flows = np.random.normal(size=flow_shape) * 3
+    return image.astype(image_type), flows.astype(flow_type)
+
+
+def _check_interpolation_correctness(
+    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 = _get_random_image_and_flows(shape, image_type, flow_type)
+
+    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])
+        y_index = np.random.randint(0, shape[1])
+        x_index = np.random.randint(0, shape[2])
+
+        _assert_correct_interpolation_value(
+            rand_image,
+            rand_flows,
+            interp,
+            batch_index,
+            y_index,
+            x_index,
+            low_precision=low_precision,
+        )
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_interpolation():
+    """Apply _check_interpolation_correctness() for a few sizes and
+    types."""
+    shapes_to_try = [[3, 4, 5, 6], [1, 2, 2, 1]]
+    for im_type in ["float32", "float64", "float16"]:
+        for flow_type in ["float32", "float64", "float16"]:
+            for shape in shapes_to_try:
+                _check_interpolation_correctness(shape, im_type, flow_type)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_size_exception():
+    """Make sure it throws an exception for images that are too small."""
+    shape = [1, 2, 1, 1]
+    errors = (ValueError, tf.errors.InvalidArgumentError)
+    with pytest.raises(errors) as exception_raised:
+        _check_interpolation_correctness(shape, "float32", "float32")
+    assert "Grid width must be at least 2." in str(exception_raised.value)
+
+
+@pytest.mark.usefixtures("maybe_run_functions_eagerly")
+def test_unknown_shapes():
+    """Apply _check_interpolation_correctness() for a few sizes and check
+    for tf.Dataset compatibility."""
+    shapes_to_try = [[3, 4, 5, 6], [1, 2, 2, 1]]
+    for shape in shapes_to_try:
+        _check_interpolation_correctness(shape, "float32", "float32", True)
 
 
 if __name__ == "__main__":