Image Processing

Contains functions for basic image processing needs

This module contains two kinds of image processing tools. The first type is specific tools for calculating commonly needed metrics (such as the centroid of an image), directly on complex-valued torch tensors. The second kind of tools perform common image manipulations on torch tensors, in such a way that it is safe to include them in automatic differentiation models.

cdtools.tools.image_processing.hann_window(im)

Applies a hann window to a 2D image to apodize it

TODO: update for pytorch

Parameters:: im (np.array) – A numpy array to apply a hann window to
Returns:: apodidzed – The image, apodized by a hann window
Return type:: np.array

cdtools.tools.image_processing.centroid(im, dims=2)

Returns the centroid of an image or a stack of images

By default, the last two dimensions are used in the calculation and the remainder of the dimensions are passed through.

Beware that the meaning of the centroid is not well defined if your image contains values less than 0

Parameters:

im (torch.Tensor) – An image or stack of images to calculate from
dims (int) – Default 2, how many trailing dimensions to calculate with

Returns:

centroid – An (i,j) index or stack of indices

Return type:

torch.Tensor

cdtools.tools.image_processing.centroid_sq(im, dims=2, comp=False)

Returns the centroid of the square of an image or stack of images

By default, the last two dimensions are used in the calculation and the remainder of the dimensions are passed through.

If the “comp” flag is set, it will be assumed that the last dimension represents the real and imaginary part of a complex number, and the centroid will be calculated for the magnitude squared of those numbers

Parameters:

im (torch.Tensor) – An image or stack of images to calculate from
dims (int) – Default 2, how many trailing dimensions to calculate for
comp (bool) – Default is False, whether the data represents complex numbers

Returns:

centroid – An (i,j) index or stack of indices

Return type:

torch.Tensor

cdtools.tools.image_processing.sinc_subpixel_shift(im, shift)

Performs a subpixel shift with sinc interpolation on the given tensor

The subpixel shift is done circularly via a multiplication with a linear phase mask in Fourier space.

Parameters:

im (torch.Tensor) – A complex-valued tensor to perform the subpixel shift on
shift (array) – A length-2 array describing the shift to perform, in pixels

Returns:

shifted_im – The subpixel shifted tensor

Return type:

torch.Tensor

cdtools.tools.image_processing.find_subpixel_shift(im1, im2, search_around=(0, 0), resolution=10)

Calculates the subpixel shift between two images by maximizing the autocorrelation

This function only searches in a 2 pixel by 2 pixel box around the specified search_around parameter. The calculation is done using the approach outlined in “Efficient subpixel image registration algorithms”, Optics Express (2008) by Manual Guizar-Sicarios et al.

Parameters:

im1 (torch.Tensor) – The first real or complex-valued torch tensor
im2 (torch.Tensor) – The second real or complex-valued torch tensor
search_around (array) – Default (0,0), the shift to search in the vicinity of
resolution (int) – Default is 10, the fraction of a pixel to calculate to

Returns:

shift – The relative shift (i,j) needed to best map im1 onto im2

Return type:

torch.Tensor

cdtools.tools.image_processing.find_pixel_shift(im1, im2)

Calculates the integer pixel shift between two images by maximizing the autocorrelation

This function simply takes the circular correlation with an FFT and returns the position of the maximum of that correlation. This corresponds to the amount that im1 would have to be shifted by to line up best with im2

Parameters:

im1 (torch.Tensor) – The first real or complex-valued torch tensor
im2 (torch.Tensor) – The second real or complex-valued torch tensor

Returns:

shift – The integer-valued shift (i,j) that best maps im1 onto im2

Return type:

torch.Tensor

cdtools.tools.image_processing.find_shift(im1, im2, resolution=10)

Calculates the shift between two images by maximizing the autocorrelation

This function starts by calculating the maximum shift to integer pixel resolution, and then searchers the nearby area to calculate a subpixel shift

Parameters:

im1 (torch.Tensor) – The first real or complex-valued torch tensor
im2 (torch.Tensor) – The second real or complex-valued torch tensor
resolution (int) – Default is 10, the fraction of a pixel to calculate to

Returns:

shift – The relative shift (i,j) needed to best map im1 onto im2

Return type:

torch.Tensor

cdtools.tools.image_processing.convolve_1d(image, kernel, dim=0, fftshift_kernel=True)

Convolves an image with a 1d kernel along a specified dimension

The convolution is a circular convolution calculated using a Fourier transform. The calculation is done so the input remains differentiable with respect to the output.

If the image has a final dimension of 2, it is assumed to be complex. Otherwise, the image is assumed to be real. The image and kernel must either both be real or both be complex.

Parameters:

image (torch.Tensor) – The image to convolve
kernel (torch.Tensor) – The 1d kernel to convolve with
dim (int) – Default 0, the dimension to convolve along
fftshift_kernel (bool) – Default True, whether to fftshift the kernel first.

Returns:

convolved_im – The convolved image

Return type:

torch.Tensor

cdtools.tools.image_processing.center(image, image_dims=2, use_power=True, iterations=4)

Automatically centers an image or stack of images

This function is designed with probes for ptychography in mind, so the default centering method is to place the centroid of the magnitude-squared of the input image on the central pixel.

# TODO I should place the centroid actually at the zero-frequency # pixel, rather than the center of the array, so this can be used in # Fourier space also.

Because it’s intended for probes, the function will sum over all extra dimensions beyond <image_dims> when calculating the centroid, and shift all the images by the same amount. It does not calculate a separate shift for each image in the stack

It also centers by using circular shifts. This means that, after calculating the centroid position and shifting by that amount, the centroid will not be perfectly centered. To counteract this, multiple iterations are run, by default 4

Parameters:

image (torch.Tensor) – The … x N x M image to center
image_dims (int) – Default 2, the number of dimensions to center along
use_power (bool) – Default True, whether to use the square of the magnitude
iterations (int) – Default 4, the number of iterations to do

Returns:

centered_im – The centered image

Return type:

torch.Tensor