imshift_batch

Contents

imshift_batch#

ptyrad.utils.image_proc.imshift_batch(img, shifts, grid)[source]#

Generates a batch of shifted images from a single input image (…, Ny,Nx) with arbitray leading dimensions.

This function shifts a complex/real-valued input image by applying phase shifts in the Fourier domain, achieving subpixel shifts in both x and y directions.

Parameters:

  • img (torch.Tensor) – The input image to be shifted. img could be either a mixed-state complex probe (pmode, Ny, Nx) complex64 tensor, or a mixed-state pseudo-complex object stack (2,omode,Nz,Ny,Nx) float32 tensor.

  • shifts (torch.Tensor) – The shifts to be applied to the image. It should be a (Nb,2) tensor and each slice as (shift_y, shift_x).

  • grid (torch.Tensor) – The k-space grid used for computing the shifts in the Fourier domain. It should be a tensor with shape=(2, Ny, Nx), where Ny and Nx are the height and width of the images, respectively. Note that the grid is normalized so the value spans from [-0.5,0.5)

Returns:

The batch of shifted images. It has an extra dimension than the input image, i.e., shape=(Nb, …, Ny, Nx),

where Nb is the number of samples in the input batch.

Return type:

shifted_img (torch.Tensor)

Note

  • The shifts are in unit of pixel. For example, a shift of (0.5, 0.5) will shift the image by half a pixel in both y and x directions, positive is down/right-ward.

  • The function utilizes the fast Fourier transform (FFT) to perform the shifting operation efficiently.

  • Make sure to convert the input image and shifts tensor to the desired device before passing them to this function.

  • The fft2 and fftshifts are all applied on the last 2 dimensions, therefore it’s only shifting along y and x directions

  • tensor[None, …] would add an extra dimension at 0, so *[None]*ndim means unwrapping a list of ndim None as [None, None, …]

  • The img is automatically broadcast to (Nb, *img.shape), so if a batch of images are passed in, each image would be shifted independently

Contents