imageRegistration.py 12.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2016-2017 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/


__authors__ = ["J. Garriga"]
__license__ = "MIT"
29
__date__ = "07/07/2021"
30
31


32
import enum
33
34
import numpy
import cv2
35
from scipy.ndimage import center_of_mass, fourier_shift
36
from skimage import registration
37
38

from silx.utils.enum import Enum as _Enum
39

40
from darfix.io import utils
41
from.autofocus import normalized_variance
42

43

44
45
@enum.unique
class ShiftApproach(_Enum):
46
47
48
49
50
51
    """
    Different shifts approaches that can be used for the shift detection and correction.
    """
    LINEAR = "linear"
    FFT = "fft"

52
53

def compute_com(data):
54
55
56
57
    """
    Compute the center of mass of a stack of images.
    First it computes the intensity of every image (by summing its values), and then
    it uses scipy ``center_of_mass`` function to find the centroid.
58

59
60
61
    :param numpy.ndarray data: stack of images.
    :returns: the vector of intensities and the com.
    """
62
    intensity = data.sum(axis=1)
63
64
65
66
    return intensity, int(center_of_mass(intensity)[0])


def diff_com(img1, img2):
67
68
69
70
71
    """
    Finds the difference between the center of mass of two images.
    This can be used to find the shift between two images with distinguished
    center of mass.
    """
72
73
74
    return numpy.array(center_of_mass(img2)) - numpy.array(center_of_mass(img1))


75
def find_shift(img1, img2, upsample_factor=1000):
76
77
    """
    Uses the function ``register_translation`` from skimage to find the shift between two images.
78

79
80
    :param array_like img1: first image.
    :param array_like img2: second image, must be same dimensionsionality as ``img1``.
81
    :param int upsample_factor: optional.
82
    """
83
    return registration.phase_cross_correlation(img1, img2, upsample_factor=upsample_factor, return_error=False)
84
85


86
def _numpy_fft_shift(img, dx, dy):
87
    """
88
    Shift an image by Fourier approach using Numpy library.
89

90
91
92
    :param array_like img: Image to shift
    :param number dx: shift in x axis
    :param number dy: shift in y axis
93
94
    :returns: ndarray
    """
95
    img = numpy.asarray(img, dtype=numpy.float32)
96
97
98
99
100
101
102
    xnum, ynum = img.shape
    """
    Create grid to apply the Fourier transformed theory for translated images.
    Each value of each matrix (Nx, Ny) correspons to the coordinate of the corresponding pixel.
    """
    Nx, Ny = numpy.meshgrid(numpy.fft.fftfreq(ynum), numpy.fft.fftfreq(xnum))
    c = numpy.fft.fft2(img) * numpy.exp(-2 * numpy.pi * 1j * (dx * Nx + dy * Ny))
103
    return numpy.fft.ifft2(c).real
104
105


106
def _opencv_fft_shift(img, dx, dy):
107
    """
108
    Shift an image by Fourier approach using opencv library.
109

110
111
112
    :param array_like img: Image to shift
    :param number dx: shift in x axis
    :param number dy: shift in y axis
113
114
115
    :returns: ndarray
    """
    img = numpy.asarray(img, dtype=numpy.float32)
116
    ynum, xnum = img.shape
117
118
    """
    Create grid to apply the Fourier transformed theory for translated images.
119
120
    Each value of each matrix (Nx, Ny) corresponds to the coordinate of the corresponding pixel
    divided by the size of the image.
121
    """
122
123
124
125
126
127
128
129
    Nx, Ny = numpy.meshgrid(numpy.fft.fftfreq(xnum), numpy.fft.fftfreq(ynum))
    dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
    dft = dft[..., 0] + 1j * dft[..., 1]
    dft = dft * numpy.exp(-2 * numpy.pi * 1j * (dx * Nx + dy * Ny))
    real = numpy.reshape(dft.real, dft.real.shape + (1,))
    imag = numpy.reshape(dft.imag, dft.imag.shape + (1,))
    dft = numpy.concatenate((real, imag), axis=2)
    return cv2.idft(dft, flags=cv2.DFT_SCALE + cv2.DFT_REAL_OUTPUT)
130
131


132
def _scipy_fft_shift(img, dx, dy):
133
134
135
136
137
138
139
140
    """
    Shift an image by Fourier approach using scipy library.

    :param array_like img: Image to shift
    :param number dx: shift in x axis
    :param number dy: shift in y axis
    :returns: ndarray
    """
141
    return numpy.fft.ifftn(fourier_shift(numpy.fft.fftn(img), [dy, dx])).real
142
143


144
def apply_shift(img, shift, shift_approach="fft"):
145
146
147
148
149
150
151
152
153
    """
    Function to apply the shift to an image.

    :param 2-dimensional array_like img: Input array, can be complex.
    :param 2-dimensional array_like shift: The shift to be applied to the image. ``shift[0]``
        refers to the y-axis, ``shift[1]`` refers to the x-axis.
    :param Union[`linear`, `fft`] shift_approach: The shift method to be used to apply the shift.
    :returns: real ndarray
    """
154
155
    shift_approach = ShiftApproach.from_value(shift_approach)
    if shift_approach == ShiftApproach.LINEAR:
156
        img = img.astype(numpy.int16)
157
158
        M = numpy.float32([[1, 0, shift[1]], [0, 1, shift[0]]])
        return cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
159
    else:
160
        return _opencv_fft_shift(img, shift[1], shift[0])
161
162


163
164
165
166
def normalize(x):
    """
    Normalizes a vector or matrix.
    """
167
    if not numpy.count_nonzero(x):
168
        return x
169
    return x / numpy.linalg.norm(x)
170
171


172
def improve_linear_shift(data, v, h, epsilon, steps, nimages=None, shift_approach="linear"):
173
    """
174
    Function to find the best shift between the images. It loops ``steps`` times,
175
176
177
178
    applying a different shift each, and trying to find the one that has the best result.

    :param array_like data: The stack of images.
    :param 2-dimensional array_like v: The vector with the direction of the shift.
179
180
    :param float epsilon: Maximum value of h
    :param int steps: Number of different tries of h.
181
182
183
184
185
186
    :param int nimages: The number of images to be used to find the best shift. It has to
        be smaller or equal as the length of the data. If it is smaller, the images used
        are chosen using `numpy.random.choice`, without replacement.
    :param Union[`linear`,`fft`] shift_approach: The shift method to be used to apply the shift.
    :returns: ndarray
    """
187
    shift_approach = ShiftApproach.from_value(shift_approach)
188
    v = numpy.asanyarray(v)
189
    iData = range(data.shape[0])
190

191
192
    if nimages:
        iData = numpy.random.choice(iData, nimages, False)
193
194

    score = {}
195
    utils.advancement_display(0, h + epsilon, "Finding shift")
196
    step = epsilon / steps
197
    for h_ in numpy.arange(0, h + epsilon, step):
198
199
        result = numpy.zeros(data[0].shape)
        for iFrame in iData:
200
            shift = h_ * v * iFrame
201
            result += apply_shift(data[iFrame], shift, shift_approach)
202
203
204

        # Compute score using normalized variance
        # TODO: add more autofocus options
205
        score[h_] = normalized_variance(result)
206
        utils.advancement_display(h_ + step, h + epsilon, "Finding shift")
207
208
209
210
    optimal_h = max(score.keys(), key=(lambda k: score[k]))
    return optimal_h


211
def shift_detection(data, steps, shift_approach="linear"):
212
213
    """
    Finds the linear shift from a set of images.
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
214

215
216
217
218
219
    :param ndarray data: Array with the images.
    :returns: A vector of length the number of images with the linear shift
        to apply to every image.
    :rtype: ndarray
    """
220
221
    if len(data) == 0:
        return []
222
223
224
    shape = data[0].shape
    first_sum = numpy.zeros(shape)
    second_sum = numpy.zeros(shape)
225
226
    for i in range(len(data)):
        if i < len(data) / 2:
227
            first_sum += data[i]
228
        else:
229
            second_sum += data[i]
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
230
231
    shift = find_shift(first_sum, second_sum, 1000)
    v = normalize(shift)
232
233
    v_ = 2 * shift / len(data)
    h = numpy.sqrt(v_[0]**2 + v_[1]**2)
234
235
    if not h:
        return numpy.outer([0, 0], numpy.arange(len(data)))
236
    epsilon = 2 * h
237
    h = improve_linear_shift(data, v, h, epsilon, steps, shift_approach=shift_approach)
238
    return numpy.outer(h * v, numpy.arange(len(data)))
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
239

240

241
242
243
244
def sum_images(image1, image2=None):
    if image2 is None:
        return numpy.asarray(image1)
    return numpy.asarray(image1) + numpy.asarray(image2)
Julia Garriga Ferrer's avatar
Julia Garriga Ferrer committed
245

246

247
def shift_correction(data, n_shift, shift_approach="fft", callback=None):
248
249
250
    """
    Function to apply shift correction technique to stack of images.

251
252
253
254
255
256
257
258
259
    :param array_like data: The stack of images to apply the shift to.
    :param array_like n_shift: Array with the shift to be applied at every image.
        The first row has the shifts in the y-axis and the second row the shifts
        in the x-axis. For image ``i`` the shift applied will be:
        ``shift_y = n_shift[0][i] shift_x = n_shift[1][i]```.
    :param Union[`linear`,`fft`] shift_approach: Name of the shift approach
        to be used. Default: `fft`.
    :param Union[None,Function] callback: Callback function to update the
        progress.
260
261
    :returns: The shifted images.
    :rtype: ndarray
262
    """
263
    shift_approach = ShiftApproach.from_value(shift_approach)
264

265
    shift = numpy.asanyarray(n_shift)
266
267
268
269

    if not numpy.any(shift):
        return data

270
271
    assert shift.shape == (2, len(data)), "n_shift list has to be of same size as stack of images"

272
273
    shifted_data = numpy.empty(data.shape, dtype=numpy.float32)

274
    utils.advancement_display(0, len(data), "Applying shift")
275
276

    for count, frame in enumerate(data):
277
278
        shifted_data[count] = apply_shift(frame, shift[:, count], shift_approach)
        utils.advancement_display(count + 1, len(data), "Applying shift")
279
280
281
282
283
284
285
286
287
        if callback:
            callback(int(count / len(data) * 100))

    if data.dtype.kind == 'f':
        if data.dtype.itemsize > shifted_data.dtype.itemsize:
            import warnings
            warnings.warn("Precision loss to float32")
    else:
        import warnings
288
        warnings.warn("Data cast to float32")
289
290
291
292

    return shifted_data


293
def random_search(data, optimal_shift, iterations, sigma=None, shift_approach='linear'):
294
295
    """
    Function that performs random search to a set of images to find an improved vector of shifts (one
296
297
298
299
300
301
302
303
304
305
306
    for each image). For this, it adds to the optimal shift a series of random samples obtained from
    a multivariate normal distribution, and selects the one with best score.

    :param array_like data: Set of images.
    :param ndarray optimal_shift: Array with two vectors with the linear shift found for axis y and x.
    :param int iterations: Number of times the random search should be done.
    :param number sigma: Standard deviation of the distribution.
    :param Union[`linear`,`fft`] shift_approach: Name of the shift approach
        to be used. Default: `linear`.
    :returns: A vector of length the number of images with the best shift found for every image.
    :rtype: ndarray
307
308
309
    """
    best_score = -numpy.inf
    best_result = numpy.empty(optimal_shift.shape)
310

311
    if sigma is None:
312
        # Using shift from second image
313
        sigma = abs(optimal_shift[:, 1] / 3)
314

315
    utils.advancement_display(0, iterations)
316
    for i in range(iterations):
317
        result = []
318
319
        normal = numpy.random.multivariate_normal((0, 0), sigma * numpy.eye(2), len(data)).T
        n_shift = optimal_shift + normal
320
        for iFrame in range(len(data)):
321
322
            shifty = n_shift[0][iFrame]
            shiftx = n_shift[1][iFrame]
323
            result += [apply_shift(data[iFrame], [shifty, shiftx], shift_approach)]
324
        score = normalized_variance(result)
325
        if best_score < score:
326
            best_result = n_shift
327
328
329
330
331
            best_score = score

        utils.advancement_display(i + 1, iterations)

    return best_result