Coverage for local_installation_linux/mumott/pipelines/utilities/image_processing.py: 91%
94 statements
« prev ^ index » next coverage.py v7.3.2, created at 2025-05-05 21:21 +0000
1"""
2Image processing methods that are necessary but not exclusive to alignment procedures.
3"""
5from typing import Tuple
6import numpy as np
7import scipy as sp
10def compute_tukey_window(
11 width: int,
12 length: int = 0,
13) -> np.ndarray[float]:
14 """
15 Tukey window, an array with decreasing values at the border from 1 to 0.
16 Here, the shape parameter $\alpha$ is always 0.2; see
17 [here](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.windows.tukey.html)
18 for more information.
20 Parameters
21 ----------
22 width
23 Width of the window.
24 length
25 Length of the window; the default is 0.
27 Returns
28 -------
29 Tukey window, an array with dimensions :attr:`width` and :attr:`length`.
30 """
32 W = sp.signal.windows.tukey(width, 0.2)
33 W = np.transpose(W[:, np.newaxis])
35 if length > 0: 35 ↛ 41line 35 didn't jump to line 41, because the condition on line 35 was never false
36 V = sp.signal.windows.tukey(length, 0.2)
37 V = V[:, np.newaxis]
38 W = V * W
39 W = W[..., np.newaxis]
40 else:
41 W = W[..., np.newaxis]
42 return W
45def get_img_grad(
46 img: np.ndarray[float],
47 x_axis_index: int = 1,
48 y_axis_index: int = 0,
49) -> Tuple[np.ndarray[float], np.ndarray[float]]:
50 """
51 Compute the vertical and horizontal gradient of the image.
53 Parameters
54 ----------
55 img
56 Image of wich we want to compute the gradient.
58 Returns
59 -------
60 Tuple of two arrays that respectively contain the horizontal
61 and vertical gradient of the image.
62 """
64 # check if its real, to enforce its type at the end
65 is_real = ~np.iscomplexobj(img)
67 Np = np.shape(img)
69 # defining a basis vector
70 vector = np.linspace(0, (Np[1] - 1) / (Np[1] + np.finfo(float).eps), Np[1])
71 Y = 1j * 2 * np.pi * (np.fft.fftshift(vector) - 0.5)
72 # Fourier transform on a list of vectors, i.e., FFT in one direction only
73 f_img_Y = np.fft.fft(img, axis=1)
74 # compute gradient in Fourier space
75 dY = f_img_Y * np.transpose(Y[:, np.newaxis])
76 # back to direct space
77 dY = np.fft.ifft(dY, axis=1)
79 vector = np.linspace(0, (Np[0] - 1) / (Np[0] + np.finfo(float).eps), Np[0])
80 X = 1j * 2 * np.pi * (np.fft.fftshift(vector) - 0.5)
81 # Fourier transform on a list of vector, i.e., FFT in one direction only
82 f_img_X = np.fft.fft(img, axis=0)
83 # compute gradient in Fourier space
84 dX = np.transpose(np.transpose(f_img_X) * X[np.newaxis, :])
85 # back to direct space
86 dX = np.fft.ifft(dX, axis=0)
88 # force real if it was
89 if is_real: 89 ↛ 93line 89 didn't jump to line 93, because the condition on line 89 was never false
90 dX = np.real(dX)
91 dY = np.real(dY)
93 if x_axis_index == 0:
94 return dX, dY
95 # else invert them
96 return dY, dX
99def imshift_fft(
100 img: np.ndarray[float],
101 x: np.ndarray[float],
102 y: np.ndarray[float],
103) -> np.ndarray[float]:
104 """
105 Apply subpixel shift to a stack of images in direct space.
107 Parameters
108 ----------
109 img
110 Stack of images to translate; the array should have three dimensions,
111 the last of which corresponds to the index relative to the stack.
112 x
113 Translation to apply along the 0 axis, for each image of the stack.
114 y
115 Translation to apply along the 1 axis, for each image of the stack.
117 Returns
118 -------
119 Translated stack of images as a three-dimensional array.
120 """
122 if len(x) != img.shape[-1]: 122 ↛ 123line 122 didn't jump to line 123, because the condition on line 122 was never true
123 raise ValueError(f'The length of x ({len(x)}) does not match the'
124 f' third dimension of the input stack ({img.shape[-1]}).')
125 if len(y) != img.shape[-1]: 125 ↛ 126line 125 didn't jump to line 126, because the condition on line 125 was never true
126 raise ValueError(f'The length of y ({len(y)}) does not match the'
127 f' third dimension of the input stack ({img.shape[-1]}).')
129 img_t = np.zeros(img.shape, dtype=complex)
131 # apply shift per image
132 for index in range(img.shape[-1]):
133 # do not use FFT shift, as its in a way included in the fourier_shift function,
134 # and back to direct space
135 tmp = np.fft.fft2(img[:, :, index])
136 tmp = sp.ndimage.fourier_shift(tmp, (x[index], y[index]))
137 img_t[:, :, index] = np.fft.ifft2(tmp)
138 return np.real(img_t)
141def smooth_edges(
142 img: np.ndarray[float],
143 smooth_kernel: int = 3,
144) -> np.ndarray[float]:
145 """Given a stack of 2D images this function smoothens the border of
146 the images to avoid sharp edge artefacts during imshift_fft.
148 Parameters
149 ----------
150 img
151 Stack of images, for which to smoothen the edges; the array should
152 have three dimensions, the last of which corresponds to the index
153 relative to the stack.
154 smooth_kernel : int, optional
155 Size of the smoothing region.
157 Returns
158 -------
159 Stack of images with smoothened edges.
160 """
162 # init
163 Npix = np.shape(img)
164 smooth_img = img
166 # we do the smoothing on the two axes
167 for i in range(2):
168 # the middle coordinates of the image along the current axis
169 half = int(np.ceil(Npix[i] / 2))
170 # roll push and we want to pull, therefore we put a negative sign
171 smooth_img = np.roll(smooth_img, -half, axis=i)
172 # define the range of data that we are going to smoothen
173 indices = half + np.linspace(-smooth_kernel, smooth_kernel, 2 * smooth_kernel + 1).astype(int) - 1
174 # define the smoothing kernel
175 ker_size = [1, 1, 1]
176 # on the appropriate dimension we shape the kernel to the wanted size
177 ker_size[i] = smooth_kernel
179 # we extract the data that we want to smoothen, on the appropriate range and axis
180 if i == 0:
181 img_tmp = smooth_img[indices, :, ...]
182 else:
183 img_tmp = smooth_img[:, indices, ...]
184 # smoothen across the image edges by convolution
185 img_tmp = sp.ndimage.convolve(img_tmp, np.ones(ker_size), mode='constant')
187 # avoid boundary issues
188 boundary_shape = [1, 1, 1]
189 boundary_shape[i] = 2 * smooth_kernel + 1
191 # compute the core convolution that was applied
192 convolution = sp.ndimage.convolve(np.ones(boundary_shape), np.ones(ker_size), mode='constant')
194 # remove it from the image
195 img_tmp = img_tmp / (convolution + np.finfo(float).eps)
196 # what is left is just a convolution of the border
198 # store the result on top of the original image, on the appriate range and axis
199 if i == 0:
200 smooth_img[indices, :, ...] = img_tmp
201 else:
202 smooth_img[:, indices, ...] = img_tmp
203 # roll push
204 smooth_img = np.roll(smooth_img, half, axis=i)
205 return smooth_img
208def imfilter_high_pass_1d(
209 img: np.ndarray[float],
210 ax: int,
211 sigma: float,
212) -> np.ndarray[float]:
213 """Applies an FFT filter along the :attr:`ax` dimension that removes
214 :attr:`sigma` ratio of the low frequencies.
216 Parameters
217 ----------
218 img
219 Input image as a two-dimensional array.
220 ax
221 Filtering axis.
222 sigma
223 Filtering intensity; should be between 0 and 1,
224 where a value <=0 implies no filtering.
226 Returns
227 -------
228 The filtered image.
229 """
231 if sigma <= 0: 231 ↛ 232line 231 didn't jump to line 232, because the condition on line 231 was never true
232 return img
233 Ndims = np.ndim(img)
234 Npix = np.shape(img)
235 shape = np.ones((Ndims))
236 shape[ax] = Npix[ax]
238 isReal = ~np.iscomplexobj(img)
240 img_f = np.fft.fft(img, axis=ax)
242 # create
243 x = np.linspace(-Npix[ax] / 2, Npix[ax] / 2 - 1, Npix[ax]) / (Npix[ax] + np.finfo(float).eps)
244 x = np.reshape(x, shape.astype(int))
246 # compute spectral filter
247 spectral_filter = np.fft.fftshift(
248 np.exp(1 / (-(x**2 + np.finfo(float).eps) / (sigma**2 + np.finfo(float).eps)))
249 )
251 img_f = img_f * spectral_filter
252 img_filt = np.fft.ifft(img_f, axis=ax)
254 if isReal: 254 ↛ 256line 254 didn't jump to line 256, because the condition on line 254 was never false
255 img_filt = np.real(img_filt)
256 return img_filt
259def center(
260 tomogram: np.ndarray[float],
261 use_shift: bool = True,
262) -> Tuple[float, float, float]:
263 """Find the center of mass of :attr:`tomogram` and return the
264 displacement vector to find it at a given rate step, calculate
265 variance if needed.
267 Parameters
268 ----------
269 tomogram
270 A tomogram in the form of a stack of 2D images; the array should have three dimensions,
271 the last of which corresponds to the index relative to the stack.
272 use_shift
273 If ``True``, the center of mass is calculated relative to the center of the image.
275 Returns
276 -------
277 A tuple that comprises the the x and y coordinates of
278 the barycenter of mass as well as the mass.
279 """
281 # size
282 size_y = np.size(tomogram, 1)
283 size_x = np.size(tomogram, 0)
285 # the mass is just the sum of the values
286 mass = np.sum(tomogram, axis=(0, 1)) + np.finfo(float).eps # epsilon to avoid zero division latter
288 # defining the grid to compute the barycenter
289 xgrid = np.arange(0, size_x) + 1
290 ygrid = np.transpose(np.arange(0, size_y) + 1)
292 # compute the barycenter of mass
293 pos_x = np.sum(xgrid[..., np.newaxis] * np.sum(tomogram, 1), 0) / mass # sum on y, to weight the x
294 pos_y = np.sum(ygrid[..., np.newaxis] * np.sum(tomogram, 0), 0) / mass # sum on x, to weigth the y
296 if use_shift: 296 ↛ 299line 296 didn't jump to line 299, because the condition on line 296 was never false
297 pos_x = pos_x - (size_x + 1) / 2 # defined so that center(ones(N), use_shift = true) == [0,0]
298 pos_y = pos_y - (size_y + 1) / 2 # defined so that center(ones(N), use_shift = true) == [0,0]
299 return pos_x, pos_y, mass # pos_y => shift on the axis 1 ; pos_x => shift on the axis 0