diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index fd160d1644..20d998afe6 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -9,7 +9,7 @@ from scipy.linalg import lstsq import aspire.volume -from aspire.nufft import anufft +from aspire.nufft import anufft, nufft from aspire.numeric import fft, xp from aspire.utils import FourierRingCorrelation, anorm, crop_pad_2d, grid_2d from aspire.volume import SymmetryGroup @@ -186,6 +186,42 @@ def __init__(self, data, dtype=None): self.__array_interface__ = self._data.__array_interface__ self.__array__ = self._data + def project(self, angles): + """ + Computes the Radon Transform on an Image Stack using + Non-Uniform Fast Fourier Transforms. This method projects the + Image stack along different angles and returns the Radon + Transform. + + :param angles: A 1-D Numpy Array of angles in Radians. + This is used to compute the Radon Transform at different angles. + :return: Radon transform of the Image Stack. + :rtype: Ndarray (stack size, number of angles, image resolution) + """ + # number of points to sample on radial line in polar grid + n_points = self.resolution + original_stack = self.stack_shape + + # 2-D grid + radial_idx = np.fft.rfftfreq(n_points) * np.pi * 2 + n_real_points = len(radial_idx) + n_angles = len(angles) + + pts = np.empty((2, n_angles, n_real_points), dtype=self.dtype) + pts[0] = radial_idx[np.newaxis, :] * np.sin(angles)[:, np.newaxis] + pts[1] = radial_idx[np.newaxis, :] * np.cos(angles)[:, np.newaxis] + pts = pts.reshape(2, n_real_points * n_angles) + + # compute the polar nufft (NUFFT) + image_ft = nufft(self.stack_reshape(-1)._data, pts).reshape( + self.n_images, n_angles, n_real_points + ) + + # Radon transform, output: (stack size, angles, points) + image_rt = np.fft.fftshift(np.fft.irfft(image_ft, n=n_points, axis=-1), axes=-1) + image_rt = image_rt.reshape(*original_stack, n_angles, n_points) + return image_rt + @property def res(self): warn( diff --git a/tests/test_sinogram.py b/tests/test_sinogram.py new file mode 100644 index 0000000000..56aa6776e0 --- /dev/null +++ b/tests/test_sinogram.py @@ -0,0 +1,138 @@ +import numpy as np +import pytest +from skimage import data +from skimage.transform import radon + +from aspire.image import Image +from aspire.utils import grid_2d + +# Relative tolerance comparing line projections to scikit +# The same tolerance will be used in all scikit comparisons +SK_TOL = 0.005 + +IMG_SIZES = [ + 511, + 512, +] + +DTYPES = [ + np.float32, + np.float64, +] + +ANGLES = [ + 1, + 50, + pytest.param(90, marks=pytest.mark.expensive), + pytest.param(117, marks=pytest.mark.expensive), + pytest.param(180, marks=pytest.mark.expensive), + pytest.param(360, marks=pytest.mark.expensive), +] + + +@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module") +def dtype(request): + """ + Dtypes for image. + """ + return request.param + + +@pytest.fixture(params=IMG_SIZES, ids=lambda x: f"px={x}", scope="module") +def img_size(request): + """ + Image size. + """ + return request.param + + +@pytest.fixture(params=ANGLES, ids=lambda x: f"n_angles={x}", scope="module") +def num_ang(request): + """ + Number of angles in radon transform. + """ + return request.param + + +@pytest.fixture +def masked_image(dtype, img_size): + """ + Creates a masked image fixture using camera data from Scikit-Image. + """ + g = grid_2d(img_size, normalized=True, shifted=True) + mask = g["r"] < 1 + + image = data.camera().astype(dtype) + image = image[:img_size, :img_size] + return Image(image * mask) + + +# Image.project and compare results to skimage.radon +def test_image_project(masked_image, num_ang): + """ + Test Image.project on a single stack of images. Compares project method output with skimage project. + """ + ny = masked_image.resolution + angles = np.linspace(0, 360, num_ang, endpoint=False) + rads = angles / 180 * np.pi + s = masked_image.project(rads) + assert s.shape == (1, len(angles), ny) + + # sci-kit image `radon` reference + # + # Note, Image.project's angles are wrt projection line (ie + # grid), while sk's radon are wrt the image. To correspond the + # rotations are inverted. This was the convention prefered by + # the original author of this method. + # + # Note, transpose sk output to match (angles, points) + reference_sinogram = radon(masked_image._data[0], theta=angles[::-1]).T + assert reference_sinogram.shape == (len(angles), ny), "Incorrect Shape" + + # compare project method on ski-image reference + nrms = np.sqrt(np.mean((s[0] - reference_sinogram) ** 2, axis=-1)) / np.linalg.norm( + reference_sinogram, axis=-1 + ) + + np.testing.assert_array_less(nrms, SK_TOL, "Error in image projections.") + + +def test_multidim(num_ang): + """ + Test Image.project on stacks of images. Extension of test_image_project but for multi-dimensional stacks. + """ + + L = 512 # pixels + n = 3 + m = 2 + + # Generate a mask + g = grid_2d(L, normalized=True, shifted=True) + mask = g["r"] < 1 + + # Generate images + imgs = Image(np.random.random((m, n, L, L))) * mask + + # Generate line project angles + angles = np.linspace(0, 360, num_ang, endpoint=False) + rads = angles / 180.0 * np.pi + s = imgs.project(rads) + + # Compare + reference_sinograms = np.empty((m, n, num_ang, L)) + for i in range(m): + for j in range(n): + img = imgs[i, j] + # Compute the singleton case, and compare with stack. + single_sinogram = img.project(rads) + + # These should be allclose up to determinism in the FFT and NUFFT. + np.testing.assert_allclose(s[i, j : j + 1], single_sinogram) + + # Next individually compute sk's radon transform for each image. + reference_sinograms[i, j] = radon(img._data[0], theta=angles[::-1]).T + + _nrms = np.sqrt(np.mean((s - reference_sinograms) ** 2, axis=-1)) / np.linalg.norm( + reference_sinograms, axis=-1 + ) + np.testing.assert_array_less(_nrms, SK_TOL, "Error in image projections.")