Commit 6b7c1886 authored by Nicola Vigano's avatar Nicola Vigano
Browse files

Implement support for n-dimensional correlations (2D is default)



Signed-off-by: Nicola Vigano's avatarNicola VIGANÒ <nicola.vigano@esrf.fr>
parent 4aefb158
Pipeline #27871 passed with stages
in 3 minutes and 48 seconds
......@@ -278,23 +278,34 @@ class AlignmentBase(object):
return img
@staticmethod
def _compute_correlation_fft(img_1, img_2, padding_mode):
def _compute_correlation_fft(img_1, img_2, padding_mode, axes=(-2, -1)):
do_circular_conv = padding_mode is None or padding_mode == "wrap"
if not do_circular_conv:
padding = np.ceil(np.array(img_2.shape) / 2).astype(np.int)
img_1 = np.pad(img_1, ((padding[0],), (padding[1],)), mode=padding_mode)
img_2 = np.pad(img_2, ((padding[0],), (padding[1],)), mode=padding_mode)
img_shape = img_2.shape
pad_size = np.ceil(np.array(img_shape) / 2).astype(np.int)
pad_array = [(0,)] * len(img_shape)
for a in axes:
pad_array[a] = (pad_size[a],)
img_1 = np.pad(img_1, pad_array, mode=padding_mode)
img_2 = np.pad(img_2, pad_array, mode=padding_mode)
# compute fft's of the 2 images
img_fft_1 = np.fft.fft2(img_1)
img_fft_2 = np.conjugate(np.fft.fft2(img_2))
img_fft_1 = np.fft.fftn(img_1, axes=axes)
img_fft_2 = np.conjugate(np.fft.fftn(img_2, axes=axes))
# inverse fft of the product to get cross_correlation of the 2 images
cc = np.real(np.fft.ifft2(img_fft_1 * img_fft_2))
cc = np.real(np.fft.ifftn(img_fft_1 * img_fft_2, axes=axes))
if not do_circular_conv:
cc = np.fft.fftshift(cc, axes=(-2, -1))
cc = cc[padding[0] : -padding[0], padding[1] : -padding[1]]
cc = np.fft.ifftshift(cc, axes=(-2, -1))
cc = np.fft.fftshift(cc, axes=axes)
slicing = [slice(None)] * len(img_shape)
for a in axes:
slicing[a] = slice(pad_size[a], cc.shape[a]-pad_size[a])
print(slicing)
cc = cc[tuple(slicing)]
cc = np.fft.ifftshift(cc, axes=axes)
return cc
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment