nanopyx.core.analysis.rcc

  1# REF: based on https://github.com/jungmannlab/picasso/blob/d867f561ffeafce752f37968a20698556d04dafb/picasso/imageprocess.py
  2
  3import numpy as np
  4import lmfit
  5from tqdm import tqdm
  6
  7from .ccm import calculate_ccm_from_ref
  8from .ccm_helper_functions import check_even_square, make_even_square
  9
 10# TODO: fix max_shift parameter
 11
 12
 13def calculate_x_corr(im1: np.ndarray, im2: np.ndarray):
 14    """
 15    Calculate the cross-correlation between two images.
 16
 17    Args:
 18        im1 (np.ndarray): The first image as a NumPy array.
 19        im2 (np.ndarray): The second image as a NumPy array.
 20
 21    Returns:
 22        np.ndarray: The cross-correlation matrix as a NumPy array.
 23    """
 24    ccm = calculate_ccm_from_ref(np.array([im2]).astype(np.float32), im1.astype(np.float32))[0]
 25
 26    return np.array(ccm)
 27
 28
 29def get_image_shift(im1: np.ndarray, im2: np.ndarray, box: int, max_shift: int = None):
 30    """
 31    Calculate the shift in coordinates between two input images.
 32
 33    Parameters:
 34    - im1: A numpy array representing the first image.
 35    - im2: A numpy array representing the second image.
 36    - box: An integer specifying the size of the fitting region of interest (ROI) in pixels.
 37    - max_shift: An optional integer specifying the maximum allowed shift in pixels. If not provided, the entire image will be used for correlation.
 38
 39    Returns:
 40    - A tuple of two integers representing the shift in the x and y coordinates, respectively. The shift is calculated as the difference between the centroid of the fitting ROI and the centroid of the correlation peak.
 41
 42    Note:
 43    - If either of the input images is completely black (i.e., all pixels have a value of zero), the function will return (0, 0).
 44    - The correlation between the two images is computed using the `calculate_x_corr` function.
 45    - The correlation matrix is cropped based on the `max_shift` parameter. If the `max_shift` is positive, the correlation matrix is cropped symmetrically around the center of the image. If the `max_shift` is zero or negative, the entire correlation matrix is used.
 46    - The brightest pixel in the correlation matrix is identified, and a fitting ROI centered on this pixel is extracted. The size of the fitting ROI is specified by the `box` parameter.
 47    - The function fits a 2D Gaussian model to the values in the fitting ROI using the `lmfit` library. The parameters of the Gaussian model are then used to calculate the centroid coordinates.
 48    - The centroid coordinates are adjusted to account for the cropping of the correlation matrix and the size of the fitting ROI.
 49    - Finally, the calculated shift in coordinates is returned as a negative value to align the images.
 50
 51    Example usage:
 52    shift_x, shift_y = get_image_shift(image1, image2, 10, max_shift=5)
 53    """
 54    if (np.sum(im1) == 0) or (np.sum(im2) == 0):
 55        return 0, 0
 56
 57    # Compute image correlation
 58    x_corr = calculate_x_corr(im1, im2)
 59
 60    # crop XCorr based on max_shift
 61    w, h = im1.shape
 62    if max_shift > 0:
 63        x_border = int((w - max_shift) / 2)
 64        y_border = int((h - max_shift) / 2)
 65        if x_border > 0:
 66            x_corr = x_corr[x_border:-x_border, :]
 67        else:
 68            x_border = 0
 69        if y_border > 0:
 70            x_corr = x_corr[:, y_border:-y_border]
 71        else:
 72            y_border = 0
 73    else:
 74        x_border = y_border = 0
 75
 76    # A quarter of the fit ROI
 77    fit_box = int(box / 2)
 78
 79    # A coordinate grid for the fitting ROI
 80    x, y = np.mgrid[-fit_box : fit_box + 1, -fit_box : fit_box + 1]
 81
 82    # Find the brightest pixel and cut out the fit ROI
 83    x_max_xc, y_max_xc = np.unravel_index(x_corr.argmax(), x_corr.shape)
 84    fit_roi = x_corr[
 85        x_max_xc - fit_box : y_max_xc + fit_box + 1,
 86        y_max_xc - fit_box : y_max_xc + fit_box + 1,
 87    ]
 88
 89    dimensions = fit_roi.shape
 90
 91    if 0 in dimensions or dimensions[0] != dimensions[1]:
 92        xc, yc = 0, 0
 93    else:
 94        # The fit model
 95        def flat_2d_gaussian(a, xc, yc, s, b):
 96            A = a * np.exp(-0.5 * ((x - xc) ** 2 + (y - yc) ** 2) / s**2) + b
 97            return A.flatten()
 98
 99        gaussian2d = lmfit.Model(flat_2d_gaussian, name="2D Gaussian", independent_vars=[])
100
101        # Set up initial parameters and fit
102        params = lmfit.Parameters()
103        params.add("a", value=fit_roi.max(), vary=True, min=0)
104        params.add("xc", value=0, vary=True)
105        params.add("yc", value=0, vary=True)
106        params.add("s", value=1, vary=True, min=0)
107        params.add("b", value=fit_roi.min(), vary=True, min=0)
108        tmp = fit_roi.flatten()
109        results = gaussian2d.fit(tmp, params)
110
111        # Get maximum coordinates and add offsets
112        xc = results.best_values["xc"]
113        yc = results.best_values["yc"]
114        xc += y_border + x_max_xc
115        yc += x_border + y_max_xc
116
117        xc -= np.floor(w / 2)
118        yc -= np.floor(h / 2)
119
120    return -xc, -yc
121
122
123def rcc(im_frames: np.ndarray, max_shift=None) -> tuple:
124    """
125    Calculate the relative shifts between a set of input images using the RCC (Relative Cross Correlation) algorithm.
126
127    Parameters:
128    - im_frames: A NumPy array representing a set of input images. The shape of the array is (n_frames, height, width).
129    - max_shift: An optional parameter specifying the maximum shift allowed between image pairs. If not provided, the default value is None.
130
131    Returns:
132    - A tuple containing the minimized shifts in the x and y directions between all image pairs.
133
134    References:
135    - REF: https://github.com/yinawang28/RCC
136    """
137    if not check_even_square(im_frames.astype(np.float32)):
138        im_frames = np.array(make_even_square(im_frames.astype(np.float32)))
139    n_frames = im_frames.shape[0]
140    shifts_x = np.zeros((n_frames, n_frames))
141    shifts_y = np.zeros((n_frames, n_frames))
142    n_pairs = int(n_frames * (n_frames - 1) / 2)
143    flag = 0
144
145    with tqdm(total=n_pairs, desc="Correlating image pairs", unit="pairs") as progress_bar:
146        for i in range(n_frames - 1):
147            for j in range(i + 1, n_frames):
148                progress_bar.update()
149                shifts_x[i, j], shifts_y[i, j] = get_image_shift(im_frames[i], im_frames[j], 5, max_shift)
150                flag += 1
151
152    return minimize_shifts(shifts_x, shifts_y)
153
154
155def minimize_shifts(shifts_x, shifts_y, shifts_z=None):
156    """
157    Functions that minimizes from the input arrays.
158
159    Parameters:
160        shifts_x (ndarray): An array containing the shifts in the x direction.
161        shifts_y (ndarray): An array containing the shifts in the y direction.
162        shifts_z (ndarray, optional): An array containing the shifts in the z direction. Defaults to None.
163
164    Returns:
165        tuple: A tuple containing the shift values in the x, y, and z directions (if shifts_z is not None) or just the shift values in the x and y directions.
166    """
167    n_channels = shifts_x.shape[0]
168    n_pairs = int(n_channels * (n_channels - 1) / 2)
169    n_dims = 2 if shifts_z is None else 3
170    rij = np.zeros((n_pairs, n_dims))
171    A = np.zeros((n_pairs, n_channels - 1))
172    flag = 0
173    for i in range(n_channels - 1):
174        for j in range(i + 1, n_channels):
175            rij[flag, 0] = shifts_x[i, j]
176            rij[flag, 1] = shifts_y[i, j]
177            if n_dims == 3:
178                rij[flag, 2] = shifts_z[i, j]
179            A[flag, i:j] = 1
180            flag += 1
181    Dj = np.dot(np.linalg.pinv(A), rij)
182    shift_x = np.insert(np.cumsum(Dj[:, 0]), 0, 0)
183    shift_y = np.insert(np.cumsum(Dj[:, 1]), 0, 0)
184    if n_dims == 2:
185        return shift_x, shift_y
186    else:
187        shift_z = np.insert(np.cumsum(Dj[:, 2]), 0, 0)
188        return shift_x, shift_y, shift_z