nanopyx.methods.channel_registration.estimator

  1import numpy as np
  2from numpy import array
  3from skimage.io import imsave
  4
  5from .corrector import ChannelRegistrationCorrector
  6from ...core.analysis._le_channel_registration import (
  7    ChannelRegistrationEstimator as leChannelRegistrationEstimator,
  8)
  9
 10
 11# this class assumes that the image is a numpy array with shape = [n_channels, height, width]
 12# assumes that channels in an image that will be aligned using generated translation masks will be in the same order
 13class ChannelRegistrationEstimator(object):
 14    """
 15    A class for estimating and applying channel registration to an image stack.
 16
 17    This class assumes that the image is a numpy array with shape [n_channels, height, width].
 18    It also assumes that the channels in an image that will be aligned using generated translation masks
 19    will be in the same order.
 20
 21    Attributes:
 22        translation_masks (numpy.ndarray): An array to store translation masks.
 23        ccms (numpy.ndarray): An array to store cross-correlation matrices.
 24
 25    Methods:
 26        apply_elastic_transform(img_stack): Apply elastic transformations to align channels.
 27        calculate_translation(channel_to_align, ref_channel_img, max_shift, blocks_per_axis, min_similarity, algorithm, method): Calculate translation masks for channel alignment.
 28        save_translation_mask(path): Save translation masks to a file.
 29        estimate(img_stack, ref_channel, max_shift, blocks_per_axis, min_similarity, method, save_translation_masks, translation_mask_save_path, algorithm, apply): Estimate and apply channel registration.
 30
 31    Note:
 32        This class is designed for estimating and applying channel registration to an image stack.
 33        It calculates translation masks for channel alignment and provides options for saving results to files.
 34    """
 35
 36    def __init__(self, verbose=True) -> None:
 37        """
 38        Initialize a ChannelRegistrationEstimator instance.
 39        """
 40        self.verbose = verbose
 41        self.translation_masks = None
 42
 43    def apply_elastic_transform(self, img_stack):
 44        """
 45        Apply elastic transformations to align channels in an image stack.
 46
 47        Args:
 48            img_stack (numpy.ndarray): The image stack with shape [n_channels, height, width].
 49
 50        Returns:
 51            numpy.ndarray: The aligned image stack.
 52
 53        Example:
 54            estimator = ChannelRegistrationEstimator()
 55            aligned_stack = estimator.apply_elastic_transform(img_stack)
 56
 57        Note:
 58            This method uses the ChannelRegistrationCorrector to apply elastic transformations
 59            for aligning the channels in the provided image stack based on the stored translation masks.
 60        """
 61        corrector = ChannelRegistrationCorrector()
 62        return corrector.align_channels(
 63            img_stack, translation_masks=self.translation_masks
 64        )
 65
 66    def save_translation_mask(self, path=None):
 67        """
 68        Save the translation masks to a file.
 69
 70        Args:
 71            path (str, optional): The file path to save the translation masks.
 72                If not provided, a user input prompt will be used to specify the path.
 73                The default file name will be "_translation_masks.tif" appended to the specified path.
 74
 75        Example:
 76            estimator = ChannelRegistrationEstimator()
 77            estimator.save_translation_mask("translation_masks.tif")
 78
 79        Note:
 80            This method saves the translation masks to a TIFF file. If the `path` argument is not provided,
 81            it prompts the user to input a file path and appends "_translation_masks.tif" to it as the default file name.
 82        """
 83        if path is None:
 84            path = (
 85                input(
 86                    "Please provide a filepath to save the translation masks"
 87                )
 88                + "_translation_masks.tif"
 89            )
 90
 91        imsave(path + "_translation_masks.tif", self.translation_masks)
 92
 93    def estimate(
 94        self,
 95        img_stack: array,
 96        ref_channel: int,
 97        max_shift: float,
 98        blocks_per_axis: int,
 99        min_similarity: float,
100        save_translation_masks: bool = True,
101        translation_mask_save_path: str = None,
102        apply: bool = False,
103    ):
104        """
105        Estimate and perform channel registration on an image stack.
106
107        Args:
108            img_stack (numpy.ndarray): The image stack with shape [n_channels, height, width].
109            ref_channel (int): The reference channel index for alignment.
110            max_shift (float): Maximum shift allowed for alignment.
111            blocks_per_axis (int): Number of blocks per axis for cross-correlation.
112            min_similarity (float): Minimum similarity threshold for alignment.
113            save_translation_masks (bool, optional): Whether to save translation masks (default is True).
114            translation_mask_save_path (str, optional): The file path to save translation masks.
115                If not provided, a user input prompt will be used to specify the path.
116            apply (bool, optional): Whether to apply elastic transformation to the image stack (default is False).
117
118        Returns:
119            numpy.ndarray or None: If `apply` is True, returns the aligned image stack; otherwise, returns None.
120
121        Example:
122            estimator = ChannelRegistrationEstimator()
123            aligned_stack = estimator.estimate(
124                img_stack, ref_channel=0, max_shift=1.0, blocks_per_axis=32, min_similarity=0.5
125            )
126
127        Note:
128            This method estimates channel registration for aligning channels in the provided image stack.
129            It calculates translation masks and cross-correlation matrices (ccms) for alignment.
130            The alignment is performed based on the specified parameters, and the results can be optionally saved.
131            If `apply` is True, it applies elastic transformation to the image stack and returns the aligned stack.
132        """
133
134        if ref_channel > img_stack.shape[0]:
135            print(
136                "Reference channel number cannot be bigger than number of channels!"
137            )
138            return None
139
140        estimator = leChannelRegistrationEstimator(verbose=self.verbose)
141        self.translation_masks = estimator.run(
142            np.asarray(img_stack, dtype=np.float32),
143            ref_channel,
144            max_shift,
145            blocks_per_axis,
146            min_similarity,
147        )
148
149        if save_translation_masks:
150            self.save_translation_mask(path=translation_mask_save_path)
151
152        if apply:
153            return self.apply_elastic_transform(img_stack)
154        else:
155            return None