Initial commit for autocenter()

skued/__init__.py

@@ -43,6 +43,7 @@
 )
 from .image import (
     align,
+    autocenter,
     azimuthal_average,
     combine_masks,
     detector_scattvectors,

skued/image/__init__.py

@@ -3,6 +3,7 @@
 
 from .alignment import align, ialign, itrack_peak
 from .calibration import detector_scattvectors, powder_calq
+from .center import autocenter
 from .metrics import (
     combine_masks,
     isnr,

skued/image/center.py

@@ -0,0 +1,62 @@
+# -*- coding: utf-8 -*-
+"""
+Determine the center of diffraction images
+==========================================
+"""
+
+import numpy as np
+from skimage.registration import phase_cross_correlation
+from scipy.ndimage import shift
+
+
+def autocenter(im, mask=None):
+    """
+    Find the center of a diffraction pattern automatically.
+
+    Parameters
+    ----------
+    im : ndarray, shape (N, M)
+        Diffraction pattern.
+    mask : ndarray, shape (N,M), dtype bool, optional
+        Mask that evaluates to `True` on pixels that
+        should be used to determine the center.
+
+    Returns
+    -------
+    r, c : 2-tupe of ints
+        Indices of the center, such that `im[r, c]` is the center of the pattern.
+    """
+    im = np.asfarray(im)
+    im -= im.min()
+
+    if mask is None:
+        mask = np.ones_like(im, dtype=np.bool)
+
+    weights = im * mask.astype(im.dtype)
+
+    # Center of mass. This works because the intensity envelope of a
+    # diffraction pattern has a radial shape
+    rr, cc = np.indices(im.shape)
+    r_rough = np.average(rr, weights=weights)
+    c_rough = np.average(cc, weights=weights)
+
+    im_i = radial_inversion(im, center=(r_rough, c_rough), cval=0.0)
+    mask_i = radial_inversion(mask, center=(r_rough, c_rough), cval=False)
+
+    shift = phase_cross_correlation(
+        reference_image=im,
+        moving_image=im_i,
+        reference_mask=mask,
+        moving_mask=mask_i,
+    )
+
+    return np.array([r_rough, c_rough]) + shift / 2 - np.array([1/2, 1/2])
+
+
+def radial_inversion(im, center, cval):
+    arr_center = np.array(im.shape) / 2
+    shifted = shift(im, shift=arr_center - np.asarray(center))
+    shifted = shifted[::-1, ::-1]
+    return shift(
+        shifted, shift=np.asarray(center) - arr_center, mode="constant", cval=cval
+    )

skued/image/tests/test_center.py

@@ -0,0 +1,47 @@
+# -*- coding: utf-8 -*-
+
+from skued import autocenter, gaussian
+import numpy as np
+import itertools as it
+import pytest
+
+np.random.seed(23)
+
+def test_autocenter_trivial():
+    """ Test that autocenter() finds the center of perfect gaussian at (0,0) """
+    im = np.zeros(shape=(128, 128), dtype=np.float)
+    center = np.asarray(im.shape)/2
+    rows, cols = np.indices(im.shape)
+    im += gaussian([rows, cols], center=center, fwhm=center.max()/2)
+
+    assert np.allclose(autocenter(im), center)
+
+@pytest.mark.parametrize("rc", range(-10,10, 2))
+@pytest.mark.parametrize("cc", range(-10,10, 2))
+def test_autocenter_shifted(rc, cc):
+    """ Test that autocenter() finds the center of a shifted gaussian """
+    im = np.zeros(shape=(128, 128), dtype=np.float)
+    rows, cols = np.indices(im.shape)
+    center = np.array([64 + rc, 64 + cc])
+    im += gaussian([rows, cols], center=center, fwhm=20)
+
+    assert np.allclose(autocenter(im), center, atol=1)
+
+
+@pytest.mark.parametrize("rc", range(-10,10, 5))
+@pytest.mark.parametrize("cc", range(-10,10, 5))
+def test_autocenter_shifted_with_mask(rc, cc):
+    """ Test that autocenter() finds the center of a shifted gaussian, where
+    the center has been masked away. """
+    im = np.zeros(shape=(256, 256), dtype=np.float)
+    mask = np.ones_like(im, dtype=np.bool)
+    mask[0:130, 118:138] = False
+
+    rows, cols = np.indices(im.shape)
+    center = np.array([128 + rc, 128 + cc])
+    im += gaussian([rows, cols], center=center, fwhm=50)
+
+    im[np.logical_not(mask)] *= 0.8
+
+    assert np.allclose(autocenter(im, mask=mask), center, atol=1)
+