Replace some image tests with NumPy comparison tests

Author: ProGamerGov

tests/optim/helpers/numpy_image.py

@@ -0,0 +1,82 @@
+import numpy as np
+import torch
+
+
+def setup_batch(x: np.ndarray, batch: int = 1, dim: int = 3) -> np.ndarray:
+    assert batch > 0
+    x = x[None, :] if x.ndim == dim and batch == 1 else x
+    x = (
+        np.stack([np.copy(x) for b in range(batch)])
+        if x.ndim == dim and batch > 1
+        else x
+    )
+    return x
+
+
+class FFTImage(object):
+    """Parameterize an image using inverse real 2D FFT"""
+
+    def __init__(
+        self,
+        size=None,
+        channels: int = 3,
+        batch: int = 1,
+        init=None,
+    ) -> None:
+        super().__init__()
+        if init is None:
+            assert len(size) == 2
+            self.size = size
+        else:
+            assert init.ndim == 3 or init.ndim == 4
+            self.size = (
+                (init.shape[1], init.shape[2])
+                if init.ndim == 3
+                else (init.shape[2], init.shape[3])
+            )
+
+        frequencies = FFTImage.rfft2d_freqs(*self.size)
+        scale = 1.0 / np.maximum(
+            frequencies,
+            np.full_like(frequencies, 1.0 / (max(self.size[0], self.size[1]))),
+        )
+        scale = scale * ((self.size[0] * self.size[1]) ** (1 / 2))
+        spectrum_scale = scale[None, :, :, None]
+        self.spectrum_scale = spectrum_scale
+
+        if init is None:
+            coeffs_shape = (channels, self.size[0], self.size[1] // 2 + 1, 2)
+            random_coeffs = np.random.randn(
+                *coeffs_shape
+            )  # names=["C", "H_f", "W_f", "complex"]
+            fourier_coeffs = random_coeffs / 50
+        else:
+            fourier_coeffs = (
+                torch.rfft(torch.from_numpy(init), signal_ndim=2).numpy()
+                / spectrum_scale
+            )
+            fourier_coeffs = fourier_coeffs / spectrum_scale
+
+        fourier_coeffs = setup_batch(fourier_coeffs, batch, 4)
+        self.fourier_coeffs = fourier_coeffs
+
+    @staticmethod
+    def rfft2d_freqs(height: int, width: int) -> np.ndarray:
+        """Computes 2D spectrum frequencies."""
+        fy = np.fft.fftfreq(height)[:, None]
+        # on odd input dimensions we need to keep one additional frequency
+        wadd = 2 if width % 2 == 1 else 1
+        fx = np.fft.fftfreq(width)[: width // 2 + wadd]
+        return np.sqrt((fx * fx) + (fy * fy))
+
+    def set_image(self, correlated_image: np.ndarray) -> None:
+        coeffs = torch.rfft(torch.from_numpy(correlated_image), signal_ndim=2).numpy()
+        self.fourier_coeffs = coeffs / self.spectrum_scale
+
+    def forward(self) -> np.ndarray:
+        h, w = self.size
+        scaled_spectrum = self.fourier_coeffs * self.spectrum_scale
+        output = torch.irfft(torch.from_numpy(scaled_spectrum), signal_ndim=2)[
+            :, :, :h, :w
+        ]
+        return output.numpy()

tests/optim/param/test_images.py

@@ -1,29 +1,163 @@
 #!/usr/bin/env python3
 import unittest
 
+import numpy as np
 import torch
 
 from captum.optim._param.image import images
-from tests.helpers.basic import BaseTest, assertTensorAlmostEqual
+from tests.helpers.basic import (
+    BaseTest,
+    assertArraysAlmostEqual,
+    assertTensorAlmostEqual,
+)
+from tests.optim.helpers import numpy_image
+
+
+class TestSetupBatch(BaseTest):
+    def test_setup_batch_chw(self) -> None:
+        init = torch.randn(3, 4, 4)
+
+        batch_test = images.ImageParameterization()
+        tensor_wbatch = batch_test.setup_batch(init)
+        array_wbatch = numpy_image.setup_batch(init.numpy())
+
+        assertArraysAlmostEqual(tensor_wbatch.numpy(), array_wbatch)
+
+    def test_setup_batch_chwr(self) -> None:
+        init = torch.randn(3, 4, 4, 2)
+
+        batch_test = images.ImageParameterization()
+        tensor_wbatch = batch_test.setup_batch(init, dim=4)
+        array_wbatch = numpy_image.setup_batch(init.numpy(), dim=4)
+
+        assertArraysAlmostEqual(tensor_wbatch.numpy(), array_wbatch)
+
+    def test_setup_batch_init(self) -> None:
+        init = torch.randn(5, 3, 4, 4)
+
+        batch_test = images.ImageParameterization()
+        tensor_wbatch = batch_test.setup_batch(init, dim=3)
+        array_wbatch = numpy_image.setup_batch(init.numpy(), dim=3)
+
+        assertArraysAlmostEqual(tensor_wbatch.numpy(), array_wbatch)
 
 
 class TestFFTImage(BaseTest):
     def test_pytorch_fftfreq(self) -> None:
-        assertTensorAlmostEqual(
-            self,
-            images.FFTImage.pytorch_fftfreq(4, 4),
-            torch.tensor([0.0000, 0.0625, -0.1250, -0.0625]),
-            0,
+        assertArraysAlmostEqual(
+            images.FFTImage.pytorch_fftfreq(4, 4).numpy(), np.fft.fftfreq(4, 4)
         )
 
     def test_rfft2d_freqs(self) -> None:
-        assertTensorAlmostEqual(
-            self,
-            images.FFTImage.rfft2d_freqs(height=2, width=3),
-            torch.tensor([[0.0000, 0.3333, 0.3333], [0.5000, 0.6009, 0.6009]]),
-            delta=0.0002,
+        height = 2
+        width = 3
+        assertArraysAlmostEqual(
+            images.FFTImage.rfft2d_freqs(height, width).numpy(),
+            numpy_image.FFTImage.rfft2d_freqs(height, width),
         )
 
+    def test_fftimage_forward_randn_init(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping FFTImage test due to insufficient Torch version."
+            )
+        size = (224, 224)
+
+        fftimage = images.FFTImage(size=size)
+        fftimage_np = numpy_image.FFTImage(size=size)
+
+        fftimage_tensor = fftimage.forward()
+        fftimage_array = fftimage_np.forward()
+
+        self.assertEqual(fftimage_tensor.detach().numpy().shape, fftimage_array.shape)
+
+    def test_fftimage_forward_init_randn_batch(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping FFTImage test due to insufficient Torch version."
+            )
+        size = (224, 224)
+        batch = 5
+
+        fftimage = images.FFTImage(size=size, batch=batch)
+        fftimage_np = numpy_image.FFTImage(size=size, batch=batch)
+
+        fftimage_tensor = fftimage.forward()
+        fftimage_array = fftimage_np.forward()
+
+        self.assertEqual(fftimage_tensor.detach().numpy().shape, fftimage_array.shape)
+
+    def test_fftimage_forward_init_randn_channels(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping FFTImage test due to insufficient Torch version."
+            )
+        size = (224, 224)
+        channels = 4
+
+        fftimage = images.FFTImage(size=size, channels=channels)
+        fftimage_np = numpy_image.FFTImage(size=size, channels=channels)
+
+        fftimage_tensor = fftimage.forward()
+        fftimage_array = fftimage_np.forward()
+
+        self.assertEqual(fftimage_tensor.detach().numpy().shape, fftimage_array.shape)
+
+    def test_fftimage_forward_init_chw(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping FFTImage test due to insufficient Torch version."
+            )
+        size = (224, 224)
+        init_tensor = torch.randn(3, 224, 224)
+        init_array = init_tensor.numpy()
+
+        fftimage = images.FFTImage(size=size, init=init_tensor)
+        fftimage_np = numpy_image.FFTImage(size=size, init=init_array)
+
+        fftimage_tensor = fftimage.forward()
+        fftimage_array = fftimage_np.forward()
+
+        self.assertEqual(fftimage_tensor.detach().numpy().shape, fftimage_array.shape)
+        assertArraysAlmostEqual(fftimage_tensor.detach().numpy(), fftimage_array)
+
+    def test_fftimage_forward_init_bchw(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping FFTImage test due to insufficient Torch version."
+            )
+        size = (224, 224)
+        init_tensor = torch.randn(1, 3, 224, 224)
+        init_array = init_tensor.numpy()
+
+        fftimage = images.FFTImage(size=size, init=init_tensor)
+        fftimage_np = numpy_image.FFTImage(size=size, init=init_array)
+
+        fftimage_tensor = fftimage.forward()
+        fftimage_array = fftimage_np.forward()
+
+        self.assertEqual(fftimage_tensor.detach().numpy().shape, fftimage_array.shape)
+        assertArraysAlmostEqual(fftimage_tensor.detach().numpy(), fftimage_array)
+
+    def test_fftimage_forward_init_batch(self) -> None:
+        if torch.__version__ == "1.2.0":
+            raise unittest.SkipTest(
+                "Skipping FFTImage test due to insufficient Torch version."
+            )
+        size = (224, 224)
+        batch = 5
+        init_tensor = torch.randn(1, 3, 224, 224)
+        init_array = init_tensor.numpy()
+
+        fftimage = images.FFTImage(size=size, batch=batch, init=init_tensor)
+        fftimage_np = numpy_image.FFTImage(size=size, batch=batch, init=init_array)
+
+        fftimage_tensor = fftimage.forward()
+        fftimage_array = fftimage_np.forward()
+
+        self.assertEqual(fftimage_tensor.detach().numpy().shape, fftimage_array.shape)
+        assertArraysAlmostEqual(fftimage_tensor.detach().numpy(), fftimage_array)
+
 
 class TestPixelImage(BaseTest):
     def test_pixelimage_random(self) -> None: