diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing.py index 22c4a17..c086ea1 100644 --- a/CDTools/tools/image_processing.py +++ b/CDTools/tools/image_processing.py @@ -1,3 +1,60 @@ 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) + + diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py new file mode 100644 index 0000000..01859cc --- /dev/null +++ b/tests/tools/test_image_processing.py @@ -0,0 +1,40 @@ +from __future__ import division, print_function, absolute_import + + +import pytest +import numpy as np +import torch as t + +from CDTools.tools import image_processing, cmath +from scipy import ndimage + + +def test_centroid(): + # Test single im + im = t.rand((30,40)) + sp_centroid = ndimage.measurements.center_of_mass(im.numpy()) + centroid = image_processing.centroid(im) + assert t.allclose(centroid, t.Tensor(sp_centroid)) + + # Test stack o' ims + ims = t.rand((5,30,40)) + sp_centroids = [ndimage.measurements.center_of_mass(im.numpy()) + for im in ims] + centroids = image_processing.centroid(ims) + assert t.allclose(centroids, t.Tensor(sp_centroids)) + + +def test_centroid_sq(): + # Test single im + im = t.rand((30,40)) + sp_centroid = ndimage.measurements.center_of_mass(im.numpy()**2) + centroid = image_processing.centroid_sq(im) + assert t.allclose(centroid, t.Tensor(sp_centroid)) + + # Test complex with multiple ims + ims = t.rand((5,30,40,2)) + np_ims = cmath.torch_to_complex(ims) + sp_centroids = [ndimage.measurements.center_of_mass(np.abs(im)**2) + for im in np_ims] + centroids = image_processing.centroid_sq(ims, comp=True) + assert t.allclose(centroids, t.Tensor(np.array(sp_centroids)))