Write the centroid finding function

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)
+
+    

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)))