mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
from __future__ import division, print_function, absolute_import
|
||
import numpy as np
|
||
import torch as t
|
||
from CDTools.tools import cmath
|
||
|
||
|
||
|
||
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
|
||
|
||
Args:
|
||
im (t.Tensor) : An image or stack of images to calculate from
|
||
dims (int) : Default 2, how many trailing dimensions to calculate for
|
||
|
||
Returns:
|
||
t.Tensor : An (i,j) index or stack of indices
|
||
"""
|
||
indices = (t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims))
|
||
indices = t.meshgrid(*indices)
|
||
|
||
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
|
||
|
||
Args:
|
||
im (t.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:
|
||
t.Tensor : An (i,j) index or stack of indices
|
||
"""
|
||
if comp:
|
||
im_sq = cmath.cabssq(im)
|
||
else:
|
||
im_sq = im**2
|
||
|
||
return centroid(im_sq, dims=dims)
|
||
|
||
|
||
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.
|
||
|
||
Args:
|
||
im1 (t.Tensor): The first real or complex-valued torch tensor
|
||
im2 (t.Tensor): The second real or complex-valued torch tensor
|
||
search_around (array_like) : Default (0,0), the shift to search in the vicinity of
|
||
resolution (int): Default is 10, the resolution to calculate to in units of 1/n
|
||
"""
|
||
|
||
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
|
||
|
||
Args:
|
||
im1 (t.Tensor): The first real or complex-valued torch tensor
|
||
im2 (t.Tensor): The second real or complex-valued torch tensor
|
||
search_around (array_like) : Default (0,0), the shift to search in the vicinity of
|
||
resolution (int): Default is 10, the resolution to calculate to in units of 1/n
|
||
"""
|
||
pass
|
||
|
||
|
||
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
|
||
|
||
Args:
|
||
im1 (t.Tensor): The first real or complex-valued torch tensor
|
||
im2 (t.Tensor): The second real or complex-valued torch tensor
|
||
resolution (int): Default is 10, the resolution to calculate to in units of 1/n
|
||
"""
|
||
|