Alignment: implement alternative metrics for focus calibration

Signed-off-by: Nicola VIGANÒ <nicola.vigano@esrf.fr>

nabu/preproc/alignment.py

@@ -5,6 +5,7 @@ from numpy.polynomial.polynomial import Polynomial, polyval
 
 from nabu.utils import previouspow2
 from nabu.misc import fourier_filters
+from nabu.misc import utils
 
 try:
     from scipy.ndimage.filters import median_filter
@@ -911,6 +912,34 @@ class CameraFocus(CenterOfRotation):
                 "Image positions have jitter larger than 10% of the expected step."
             )
 
+    @staticmethod
+    def _gradient(x, axes):
+        d = [None] * len(axes)
+        for ii in range(len(axes)):
+            ind = -(ii + 1)
+            padding = [(0, 0)] * len(x.shape)
+            padding[ind] = (0, 1)
+            temp_x = np.pad(x, padding, mode='constant')
+            d[ind] = np.diff(temp_x, n=1, axis=ind)
+        return np.stack(d, axis=0)
+
+    @staticmethod
+    def _compute_metric_value(data, metric, axes=(-2, -1)):
+        if metric.lower() == 'std':
+            return np.std(data, axis=axes) / np.mean(data, axis=axes)
+        elif metric.lower() == 'grad':
+            grad_data = CameraFocus._gradient(data, axes=axes)
+            grad_mag = np.sqrt(np.sum(grad_data ** 2, axis=0))
+            return np.sum(grad_mag, axis=axes)
+        elif metric.lower() == 'psd':
+            f_data = np.fft.fftn(data, axes=axes)
+            az_data = utils.azimuthal_integration(np.abs(f_data), axes=axes, domain="fourier")
+            max_vals = np.max(az_data, axis=0)
+            az_data /= max_vals[None, :]
+            return np.mean(az_data, axis=-1)
+        else:
+            raise ValueError("Unknown metric function %s" % metric)
+
     def find_distance(
         self,
         img_stack: np.ndarray,
@@ -998,14 +1027,15 @@ class CameraFocus(CenterOfRotation):
             img_stack, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape, low_pass=low_pass, high_pass=high_pass,
         )
 
-        img_stds = np.std(img_stack, axis=(-2, -1)) / np.mean(img_stack, axis=(-2, -1))
+        img_resp = self._compute_metric_value(img_stack, metric=metric, axes=(-2, -1))
 
-        # assuming images are equispaced
+        # assuming images are equispaced!
+        # focus_step = np.mean(np.abs(np.diff(img_pos)))
         focus_step = (img_pos[-1] - img_pos[0]) / (num_imgs - 1)
 
         img_inds = np.arange(num_imgs)
-        (f_vals, f_pos) = self.extract_peak_regions_1d(img_stds, peak_radius=peak_fit_radius, cc_coords=img_inds)
-        focus_ind, img_std_max = self.refine_max_position_1d(f_vals, return_vertex_val=True)
+        (f_vals, f_pos) = self.extract_peak_regions_1d(img_resp, peak_radius=peak_fit_radius, cc_coords=img_inds)
+        focus_ind, img_resp_max = self.refine_max_position_1d(f_vals, return_vertex_val=True)
         focus_ind += f_pos[1, :]
 
         focus_pos = img_pos[0] + focus_step * focus_ind
@@ -1014,10 +1044,10 @@ class CameraFocus(CenterOfRotation):
         if self.verbose:
             print("Fitted focus motor position:", focus_pos, "and corresponding image position:", focus_ind)
             f, ax = plt.subplots(1, 1)
-            ax.stem(img_pos, img_stds)
-            ax.stem(focus_pos, img_std_max, linefmt='C1-', markerfmt='C1o')
-            ax.set_title("Images std")
-            plt.show(block=False)
+            ax.stem(img_pos, img_resp)
+            ax.stem(focus_pos, img_resp_max, linefmt='C1-', markerfmt='C1o')
+            ax.set_title("Images response (metric: %s)" % metric)
+            plt.show(block=True)
 
         return focus_pos, focus_ind
 
@@ -1203,14 +1233,14 @@ class CameraFocus(CenterOfRotation):
         ], dtype=np.int)
         img_stack = np.reshape(img_stack, block_stack_size)
 
-        img_stds = np.std(img_stack, axis=(-3, -1)) / np.mean(img_stack, axis=(-3, -1))
-        img_stds = np.reshape(img_stds, [num_imgs, -1]).transpose()
+        img_resp = self._compute_metric_value(img_stack, metric=metric, axes=(-3, -1))
+        img_resp = np.reshape(img_resp, [num_imgs, -1]).transpose()
 
         # assuming images are equispaced
         focus_step = (img_pos[-1] - img_pos[0]) / (num_imgs - 1)
 
         img_inds = np.arange(num_imgs)
-        (f_vals, f_pos) = self.extract_peak_regions_1d(img_stds, peak_radius=peak_fit_radius, cc_coords=img_inds)
+        (f_vals, f_pos) = self.extract_peak_regions_1d(img_resp, peak_radius=peak_fit_radius, cc_coords=img_inds)
         focus_inds = self.refine_max_position_1d(f_vals)
         focus_inds += f_pos[1, :]