mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-13 14:52:38 +02:00
move to a src directory layout
This commit is contained in:
Abe Levitan
@@ -0,0 +1,319 @@
|
||||
"""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.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch as t
|
||||
from CDTools.tools import propagators
|
||||
|
||||
__all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift',
|
||||
'find_subpixel_shift', 'find_pixel_shift', 'find_shift',
|
||||
'convolve_1d', 'fourier_upsample']
|
||||
|
||||
|
||||
def 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 : torch.Tensor
|
||||
An (i,j) index or stack of indices
|
||||
"""
|
||||
# For some reason this needs to be a list
|
||||
indices = [t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims)]
|
||||
indices = t.meshgrid(*indices, indexing='ij')
|
||||
|
||||
use_dims = [-dims+i for i in range(dims)]
|
||||
divisor = t.sum(im, dim=use_dims)
|
||||
centroids = [t.sum(index * im, dim=use_dims) / divisor
|
||||
for index in indices]
|
||||
|
||||
return t.stack(centroids,dim=-1)
|
||||
|
||||
|
||||
def 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: torch.Tensor
|
||||
An (i,j) index or stack of indices
|
||||
"""
|
||||
if comp:
|
||||
im_sq = t.abs(im)**2
|
||||
else:
|
||||
im_sq = im**2
|
||||
|
||||
return centroid(im_sq, dims=dims)
|
||||
|
||||
|
||||
def 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 : torch.Tensor
|
||||
The subpixel shifted tensor
|
||||
"""
|
||||
|
||||
i = t.arange(im.shape[0]) - im.shape[0]//2
|
||||
j = t.arange(im.shape[1]) - im.shape[1]//2
|
||||
I,J = t.meshgrid(i,j, indexing='ij')
|
||||
I = 2 * np.pi * I.to(t.float32) / im.shape[0]
|
||||
J = 2 * np.pi * J.to(t.float32) / im.shape[1]
|
||||
I = I.to(dtype=im.dtype,device=im.device)
|
||||
J = J.to(dtype=im.dtype,device=im.device)
|
||||
|
||||
fft_im = t.fft.fftshift(t.fft.fft2(im),dim=(-2,-1))
|
||||
shifted_fft_im = fft_im * t.exp(1j * (-shift[0]*I - shift[1]*J))
|
||||
return t.fft.ifft2(t.fft.ifftshift(shifted_fft_im, dim=(-2,-1)))
|
||||
|
||||
|
||||
|
||||
def 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 : torch.Tensor
|
||||
The relative shift (i,j) needed to best map im1 onto im2
|
||||
"""
|
||||
|
||||
#
|
||||
# Here's my approach, perhaps it's a little unconventional. I will first
|
||||
# calculate the phase correlation function as found in ____ (cite a paper
|
||||
# defining it). This is strongly peaked, so I can take a small window
|
||||
# of say, 10x10 pixels, and then do a sinc interpolation of that area
|
||||
# using an FFT with upsampling by a factor of resolution in reciprocal
|
||||
# space
|
||||
#
|
||||
cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
|
||||
|
||||
# Not sure if this is more or less stable than just the correlation
|
||||
# maximum - requires some testing
|
||||
cor = t.fft.ifft2(cor_fft / t.abs(cor_fft))
|
||||
|
||||
|
||||
# Now, I need to shift the array to pull out a contiguous window
|
||||
# around the correlation maximum
|
||||
try:
|
||||
search_around = search_around.cpu()
|
||||
except:
|
||||
search_around = t.as_tensor(search_around)
|
||||
|
||||
window_size = 15
|
||||
shift_zero = tuple(-search_around + t.tensor([window_size,window_size]))
|
||||
cor_window = t.roll(cor, shift_zero, dims=(-2,-1))[...,:2*window_size,:2*window_size]
|
||||
|
||||
# Now we upsample this window
|
||||
cor_window_fft = t.fft.fftshift(t.fft.fft2(cor_window),dim=(-2,-1))
|
||||
upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape) * resolution),
|
||||
dtype=cor.dtype,device=cor.device)
|
||||
|
||||
upsampled[...,:2*window_size,:2*window_size] = cor_window_fft
|
||||
upsampled = t.roll(upsampled,(-window_size,-window_size),dims=(0,1))
|
||||
upsampled = t.roll(t.abs(t.fft.ifft2(upsampled))**2,
|
||||
(-window_size*resolution,-window_size*resolution),
|
||||
dims=(0,1))
|
||||
|
||||
|
||||
# And we extract the shift from the window
|
||||
sh = t.as_tensor(upsampled.shape, device=upsampled.device)
|
||||
cormax = t.as_tensor([t.div(t.argmax(upsampled), sh[1],
|
||||
rounding_mode='floor'),
|
||||
t.argmax(upsampled) % sh[1]],
|
||||
device=upsampled.device)
|
||||
|
||||
sh_over_2 = t.div(sh,2,rounding_mode='floor')
|
||||
subpixel_shift = ((cormax + sh_over_2) % sh - sh_over_2).to(dtype=upsampled.dtype)
|
||||
|
||||
return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \
|
||||
subpixel_shift / resolution
|
||||
|
||||
|
||||
def 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 : torch.Tensor
|
||||
The integer-valued shift (i,j) that best maps im1 onto im2
|
||||
"""
|
||||
cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
|
||||
|
||||
# Not sure if this is more or less stable than just the correlation
|
||||
# maximum - requires some testing
|
||||
cor = t.abs(t.fft.ifft2(cor_fft / t.abs(cor_fft)))
|
||||
|
||||
|
||||
sh = t.as_tensor(cor.shape,device=im1.device)
|
||||
cormax = t.tensor([t.div(t.argmax(cor),sh[1],rounding_mode='floor'),
|
||||
t.argmax(cor) % sh[1]]).to(device=im1.device)
|
||||
|
||||
sh_over_2 = t.div(sh,2,rounding_mode='floor')
|
||||
return (cormax + sh_over_2) % sh - sh_over_2
|
||||
|
||||
|
||||
|
||||
def 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 : torch.Tensor
|
||||
The relative shift (i,j) needed to best map im1 onto im2
|
||||
"""
|
||||
integer_shift = find_pixel_shift(im1,im2)
|
||||
subpixel_shift = find_subpixel_shift(im1, im2, search_around=integer_shift,
|
||||
resolution=resolution)
|
||||
|
||||
return subpixel_shift
|
||||
|
||||
|
||||
def 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 : torch.Tensor
|
||||
The convolved image
|
||||
"""
|
||||
|
||||
|
||||
if fftshift_kernel:
|
||||
kernel = t.fft.ifftshift(kernel,dim=(-1,))
|
||||
|
||||
# We have to transpose the relevant dimension to -1 before using the fft,
|
||||
# which expects to operate on the final dimension
|
||||
trans_im = t.transpose(image, dim, -1)
|
||||
|
||||
# Take a correlation
|
||||
fft_im = t.fft.fft(trans_im)
|
||||
fft_kernel = t.fft.fft(kernel)
|
||||
trans_conv = t.fft.ifft(fft_im * fft_kernel)
|
||||
|
||||
conv_im = t.transpose(trans_conv, dim, -1)
|
||||
|
||||
return conv_im
|
||||
|
||||
|
||||
def fourier_upsample(ims, preserve_mean=False):
|
||||
# If preserve_mean is true, it preserves the mean pixel intensity
|
||||
# otherwise, it preserves the total summed intensity
|
||||
upsampled = t.zeros(ims.shape[:-2]+(2*ims.shape[-2],2*ims.shape[-1]),
|
||||
dtype=ims.dtype,
|
||||
device=ims.device)
|
||||
left = [ims.shape[-2]//2,ims.shape[-1]//2]
|
||||
right = [ims.shape[-2]//2+ims.shape[-2],
|
||||
ims.shape[-1]//2+ims.shape[-1]]
|
||||
|
||||
upsampled[...,left[0]:right[0],left[1]:right[1]] = propagators.far_field(ims)
|
||||
if preserve_mean:
|
||||
upsampled *= 2
|
||||
return propagators.inverse_far_field(upsampled)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user