nanopyx.methods.drift_alignment.estimator

  1import numpy as np
  2from math import sqrt
  3from scipy.interpolate import interp1d
  4
  5from .estimator_table import DriftEstimatorTable
  6from .corrector import DriftCorrector
  7from ...core.analysis.estimate_shift import GetMaxOptimizer
  8from ...core.utils.timeit import timeit
  9from ...core.analysis.ccm import calculate_ccm
 10from ...core.analysis.rcc import rcc
 11from ...core.analysis._le_drift_calculator import (
 12    DriftEstimator as leDriftEstimator,
 13)
 14
 15
 16class DriftEstimator(object):
 17    """
 18    Drift estimator class for estimating and correcting drift in image stacks.
 19
 20    This class provides methods for estimating and correcting drift in image stacks using cross-correlation.
 21
 22    Args:
 23        None
 24
 25    Attributes:
 26        estimator_table (DriftEstimatorTable): A table of parameters for drift estimation and correction.
 27        cross_correlation_map (numpy.ndarray): The cross-correlation map calculated during drift estimation.
 28        drift_xy (numpy.ndarray): The drift magnitude at each time point.
 29        drift_x (numpy.ndarray): The drift in the X direction at each time point.
 30        drift_y (numpy.ndarray): The drift in the Y direction at each time point.
 31
 32    Methods:
 33        __init__(): Initialize the `DriftEstimator` object.
 34
 35        estimate(image_array, **kwargs): Estimate and correct drift in an image stack.
 36
 37        compute_temporal_averaging(image_arr): Compute temporal averaging of image frames.
 38
 39        get_shift_from_ccm_slice(slice_index): Get the drift shift from a slice of the cross-correlation map.
 40
 41        get_shifts_from_ccm(): Get the drift shifts from the entire cross-correlation map.
 42
 43        create_drift_table(): Create a table of drift values.
 44
 45        save_drift_table(save_as_npy=True, path=None): Save the drift table to a file.
 46
 47        set_estimator_params(**kwargs): Set parameters for drift estimation and correction.
 48
 49    Example:
 50        estimator = DriftEstimator()
 51        drift_params = {
 52            "time_averaging": 2,
 53            "max_expected_drift": 5,
 54            "shift_calc_method": "rcc",
 55            "ref_option": 0,
 56            "apply": True,
 57        }
 58        drift_corrected_image = estimator.estimate(image_stack, **drift_params)
 59
 60    Note:
 61        The `DriftEstimator` class is used for estimating and correcting drift in image stacks.
 62        It provides methods for estimating drift using cross-correlation and applying drift correction to an image stack.
 63    """
 64
 65    def __init__(self, verbose=True):
 66        """
 67        Initialize the `DriftEstimator` object.
 68
 69        Args:
 70            None
 71
 72        Returns:
 73            None
 74
 75        Example:
 76            estimator = DriftEstimator()
 77        """
 78        self.verbose = verbose
 79        self.estimator_table = DriftEstimatorTable()
 80        self.cross_correlation_map = None
 81        self.drift_xy = None
 82        self.drift_x = None
 83        self.drift_y = None
 84
 85    # @timeit
 86    def estimate(self, image_array, **kwargs):
 87        """
 88        Estimate and correct drift in an image stack.
 89
 90        Args:
 91            image_array (numpy.ndarray): The input image stack with shape [n_slices, height, width].
 92            **kwargs: Keyword arguments for setting drift estimation parameters.
 93
 94        Returns:
 95            numpy.ndarray or None: The drift-corrected image stack if `apply` is True, else None.
 96
 97        Example:
 98            drift_params = {
 99                "time_averaging": 2,
100                "max_expected_drift": 5,
101                "ref_option": 0,
102                "apply": True,
103            }
104            drift_corrected_image = estimator.estimate(image_stack, **drift_params)
105
106        Note:
107            This method estimates and corrects drift in an image stack using specified parameters.
108        """
109        self.set_estimator_params(**kwargs)
110
111        n_slices = image_array.shape[0]
112
113        # x0, y0, x1, y1 correspond to the exact coordinates of the roi to be used or full image dims and should be a tuple
114        if (
115            self.estimator_table.params["use_roi"]
116            and self.estimator_table.params["roi"] is not None
117        ):  # crops image to roi
118            print(
119                self.estimator_table.params["use_roi"],
120                self.estimator_table.params["roi"],
121            )
122            x0, y0, x1, y1 = tuple(self.estimator_table.params["roi"])
123            image_arr = image_array[:, y0 : y1 + 1, x0 : x1 + 1]
124        else:
125            image_arr = image_array
126
127        estimator = leDriftEstimator(verbose=self.verbose)
128        self.estimator_table.drift_table = estimator.run(
129            np.asarray(image_arr, dtype=np.float32),
130            time_averaging=self.estimator_table.params["time_averaging"],
131            max_drift=self.estimator_table.params["max_expected_drift"],
132            ref_option=self.estimator_table.params["ref_option"],
133        )
134
135        if self.estimator_table.params["apply"]:
136            drift_corrector = DriftCorrector()
137            drift_corrector.estimator_table = self.estimator_table
138            tmp = drift_corrector.apply_correction(image_array)
139            return tmp
140        else:
141            return None
142
143    def save_drift_table(self, save_as_npy=True, path=None):
144        """
145        Save the drift table to a file.
146
147        Args:
148            save_as_npy (bool, optional): Whether to save the table as a NumPy binary file. Default is True.
149            path (str, optional): The file path to save the table. If not provided, a user input prompt will be used.
150
151        Returns:
152            None
153
154        Example:
155            self.save_drift_table(save_as_npy=True, path="drift_table.npy")
156
157        Note:
158            This method allows saving the drift table to a file in either NumPy binary or CSV format.
159        """
160        if save_as_npy:
161            self.estimator_table.export_npy(path=path)
162        else:
163            self.estimator_table.export_csv(path=path)
164
165    def set_estimator_params(self, **kwargs):
166        """
167        Set parameters for drift estimation and correction.
168
169        Args:
170            **kwargs: Keyword arguments for setting drift estimation parameters.
171
172        Returns:
173            None
174
175        Example:
176            params = {
177                "time_averaging": 2,
178                "max_expected_drift": 5,
179                "shift_calc_method": "rcc",
180                "ref_option": 0,
181                "apply": True,
182            }
183            self.set_estimator_params(**params)
184
185        Note:
186            This method allows setting parameters for drift estimation and correction.
187        """
188        self.estimator_table.set_params(**kwargs)