nanopyx.methods.drift_alignment.corrector

  1from .estimator_table import DriftEstimatorTable
  2from ...core.utils.timeit import timeit
  3from ...core.transform._le_interpolation_bicubic import ShiftAndMagnify
  4
  5import cv2
  6import numpy as np
  7from skimage.transform import EuclideanTransform, warp
  8
  9# TODO cv2 to LE
 10
 11
 12class DriftCorrector(object):
 13    """
 14    Main class for aligning timelapse images with drift.
 15
 16    This class is used for aligning timelapse images with drift correction. It requires a previously calculated drift table.
 17    The class implements the following methods:
 18    - apply_correction
 19    - load_estimator_table
 20    - _translate_slice
 21
 22    Args:
 23        None
 24
 25    Attributes:
 26        estimator_table (DriftEstimatorTable): An instance of the DriftEstimatorTable containing the drift table data.
 27        image_arr (numpy.ndarray): The timelapse image array with shape (n_slices, rows, columns).
 28
 29    Methods:
 30        __init__(): Initialize the `DriftCorrector` object.
 31
 32        _translate_slice(slice_idx): Translate an individual image slice based on the drift table.
 33
 34        apply_correction(image_array): Apply drift correction to the entire image array.
 35
 36        load_estimator_table(path=None): Load the drift table from a file.
 37
 38    Example:
 39        corrector = DriftCorrector()
 40        corrected_image = corrector.apply_correction(image_array)
 41        corrector.load_estimator_table("drift_table.csv")
 42
 43    Note:
 44        The `DriftCorrector` class is used for correcting drift in timelapse images using a precomputed drift table.
 45    """
 46
 47    def __init__(self):
 48        """
 49        Initialize the `DriftCorrector` object.
 50
 51        Args:
 52            None
 53
 54        Returns:
 55            None
 56
 57        Example:
 58            corrector = DriftCorrector()
 59        """
 60        self.estimator_table = DriftEstimatorTable()
 61        self.image_arr = None
 62
 63    def _translate_slice(self, slice_idx):
 64        """
 65        Translate an individual image slice based on the drift table.
 66
 67        Args:
 68            slice_idx (int): The index of the slice to be translated.
 69
 70        Returns:
 71            numpy.ndarray: The translated image slice.
 72
 73        Example:
 74            translated_slice = self._translate_slice(0)
 75
 76        Note:
 77            This method is used to translate individual image slices based on the drift information in the drift table.
 78        """
 79        drift_x = self.estimator_table.drift_table[slice_idx][1]
 80        drift_y = self.estimator_table.drift_table[slice_idx][2]
 81
 82        if drift_x == 0 and drift_y == 0:
 83            return self.image_arr[slice_idx]
 84        else:
 85            return cv2.warpAffine(
 86                self.image_arr[slice_idx].astype(np.float32),
 87                np.float32([[1, 0, drift_x], [0, 1, drift_y]]),
 88                self.image_arr[slice_idx].shape[:2][::-1],
 89            ).astype(self.image_arr.dtype)
 90
 91    # @timeit
 92    def apply_correction(self, image_array):
 93        """
 94        Apply drift correction to the entire image array.
 95
 96        Args:
 97            image_array (numpy.ndarray): The input image array with shape (n_slices, rows, columns).
 98
 99        Returns:
100            numpy.ndarray: The aligned image array with shape (n_slices, rows, columns).
101
102        Example:
103            corrected_image = self.apply_correction(image_array)
104
105        Note:
106            This is the main method of the `DriftCorrector` class, which applies drift correction to the entire image array.
107        """
108        if self.estimator_table.drift_table is not None:
109            self.image_arr = image_array
110            corrected_image = [self._translate_slice(i).astype(np.float32) for i in range(0, image_array.shape[0])]
111            return np.array(corrected_image)
112            # return np.array(translation.translate_array(image_array.astype(np.float32),
113            #                                             np.array(self.estimator_table.drift_table).astype(np.float32)))
114
115        else:
116            print("Missing drift calculation")
117            return None
118
119    def load_estimator_table(self, path=None):
120        """
121        Load the drift table from a file.
122
123        Args:
124            path (str, optional): The path to the drift table file (CSV or NPY format). Default is None.
125
126        Returns:
127            None
128
129        Example:
130            self.load_estimator_table("drift_table.csv")
131
132        Note:
133            This method is used to load the drift table data from a file into the `estimator_table` attribute.
134        """
135        if path is None:
136            path = input("Please provide a filepath to the drift table")
137
138        if path.split(".")[-1] == "npy":
139            self.estimator_table.import_npy(path)
140        elif path.split(".")[-1] == "csv":
141            self.estimator_table.import_csv(path)