buteo.array.convolution_kernels

  1"""### Convolutions kernels. ###"""
  2
  3# Standard library
  4from typing import Tuple
  5
  6# External
  7import numpy as np
  8from numba import jit
  9
 10
 11
 12@jit(nopython=True, nogil=True, cache=True, fastmath=True, inline='always')
 13def _simple_blur_kernel_2d_3x3() -> Tuple[np.ndarray, np.ndarray]:
 14    """Create a 2D blur kernel.
 15
 16    The kernel has the following form:
 17    ```python
 18    >>> weights = [
 19    ...     0.08422299, 0.12822174, 0.08422299,
 20    ...     0.12822174, 0.15022110, 0.12822174,
 21    ...     0.08422299, 0.12822174, 0.08422299,
 22    ... ]
 23    ```
 24    Returns
 25    -------
 26    Tuple[np.ndarray, np.ndarray]
 27        A tuple of two numpy arrays. The first array contains the (x, y) offsets
 28        of the kernel values. The second array contains the corresponding weights
 29        of each kernel value.
 30    """
 31    offsets = np.array([
 32        [ 1, -1], [ 1, 0], [ 1, 1],
 33        [ 0, -1], [ 0, 0], [ 0, 1],
 34        [-1, -1], [-1, 0], [-1, 1],
 35    ])
 36
 37    weights = np.array([
 38        0.08422299, 0.12822174, 0.08422299,
 39        0.12822174, 0.1502211 , 0.12822174,
 40        0.08422299, 0.12822174, 0.08422299,
 41    ], dtype=np.float32)
 42
 43    return offsets, weights
 44
 45
 46@jit(nopython=True, nogil=True, cache=True, fastmath=True, inline='always')
 47def _simple_unsharp_kernel_2d_3x3() -> Tuple[np.ndarray, np.ndarray]:
 48    """Create a 2D unsharp kernel.
 49
 50    Returns
 51    -------
 52    Tuple[np.ndarray, np.ndarray]
 53        A tuple of two numpy arrays. The first array contains the (x, y) offsets
 54        of the kernel values. The second array contains the corresponding weights
 55        of each kernel value.
 56    """
 57    offsets = np.array([
 58        [ 1, -1], [ 1, 0], [ 1, 1],
 59        [ 0, -1], [ 0, 0], [ 0, 1],
 60        [-1, -1], [-1, 0], [-1, 1],
 61    ])
 62
 63    weights = np.array([
 64        -0.09911165, -0.15088834, -0.09911165,
 65        -0.15088834,  2.        , -0.15088834,
 66        -0.09911165, -0.15088834, -0.09911165,
 67    ], dtype=np.float32)
 68
 69    return offsets, weights
 70
 71
 72@jit(nopython=True, nogil=True, cache=True, fastmath=True, inline='always')
 73def _simple_shift_kernel_2d(
 74    x_offset: float,
 75    y_offset: float,
 76) -> Tuple[np.ndarray, np.ndarray]:
 77    """Create a 2D shift kernel. For augmentations."""
 78    return kernel_shift(x_offset, y_offset)
 79
 80
 81@jit(nopython=True, nogil=True, cache=True, fastmath=True, inline="always")
 82def _distance_2D(p1: np.ndarray, p2: np.ndarray) -> float:
 83    """Returns the distance between two points. (2D)"""
 84    d1 = (p1[0] - p2[0]) ** 2
 85    d2 = (p1[1] - p2[1]) ** 2
 86
 87    return np.sqrt(d1 + d2)
 88
 89
 90@jit(nopython=True, nogil=True, cache=True, fastmath=True, parallel=True)
 91def _area_covered(square, radius):
 92    """Calculates the area covered by a circle within a square.
 93    Monte-carlo(ish) method. Can be parallelized.
 94    """
 95    n_points = 100
 96    min_y = square[:, 0].min()
 97    max_y = square[:, 0].max()
 98    min_x = square[:, 1].min()
 99    max_x = square[:, 1].max()
100
101    steps = int(np.rint(np.sqrt(n_points)))
102    range_y = np.linspace(min_x, max_x, steps)
103    range_x = np.linspace(min_y, max_y, steps)
104
105    center = np.array([0.0, 0.0], dtype=np.float32)
106    adjusted_radius = radius + 0.5
107
108    points_within = 0
109    for y in range_y:
110        for x in range_x:
111            point = np.array([y, x], dtype=np.float32)
112            if _distance_2D(center, point) <= adjusted_radius:
113                points_within += 1
114
115    area = points_within / (steps ** 2)
116
117    return area
118
119
120@jit(nopython=True, nogil=True, cache=True, fastmath=True, parallel=True)
121def _circular_kernel_2D(radius):
122    """Creates a circular 2D kernel. Supports fractional radii."""
123    size = np.int64(np.ceil(radius) * 2 + 1)
124    kernel = np.zeros((size, size), dtype=np.float32)
125
126    center = np.array([0.0, 0.0], dtype=np.float32)
127
128    step = size // 2
129    for idx_i, col in enumerate(range(-step, step + 1)):
130        for idx_j, row in enumerate(range(-step, step + 1)):
131            square = np.zeros((4, 2), dtype=np.float32)
132            square[0] = np.array([col - 0.5, row - 0.5], dtype=np.float32)
133            square[1] = np.array([col + 0.5, row - 0.5], dtype=np.float32)
134            square[2] = np.array([col + 0.5, row + 0.5], dtype=np.float32)
135            square[3] = np.array([col - 0.5, row + 0.5], dtype=np.float32)
136
137            within = np.zeros(4, dtype=np.uint8)
138            for i in range(4):
139                within[i] = _distance_2D(center, square[i]) <= radius + 0.5
140
141            # Case 1: completely inside
142            if within.sum() == 4:
143                kernel[idx_i][idx_j] = 1.0
144
145            # Case 2: completely outside
146            elif within.sum() == 0:
147                kernel[idx_i][idx_j] = 0.0
148
149            # Case 3: partially inside
150            else:
151                kernel[idx_i][idx_j] = _area_covered(square, radius)
152
153    return kernel
154
155
156@jit(nopython=True, nogil=True, cache=True, fastmath=True, parallel=True)
157def _distance_weighted_kernel_2D(radius, method, decay=0.2, sigma=2.0):
158    """Creates a distance weighted kernel.
159
160    Parameters
161    ----------
162
163    radius : float
164        Radius of the kernel.
165
166    method : int
167        Method to use for weighting.
168        0. linear
169        1. sqrt
170        2. power
171        3. log
172        4. gaussian
173        5. constant
174    """
175    size = np.int64(np.ceil(radius) * 2 + 1)
176    kernel = np.zeros((size, size), dtype=np.float32)
177
178    center = np.array([0.0, 0.0], dtype=np.float32)
179
180    step = size // 2
181    for idx_i, col in enumerate(range(-step, step + 1)):
182        for idx_j, row in enumerate(range(-step, step + 1)):
183            point = np.array([col, row], dtype=np.float32)
184            distance = _distance_2D(center, point)
185
186            # Linear
187            if method == 0:
188                kernel[idx_i, idx_j] = np.power((1 - decay), distance)
189
190            # Sqrt
191            elif method == 1:
192                kernel[idx_i, idx_j] = np.power(np.sqrt((1 - decay)), distance)
193
194            # Power
195            elif method == 2:
196                kernel[idx_i, idx_j] = np.power(np.power((1 - decay), 2), distance)
197
198            # Gaussian
199            elif method == 3:
200                kernel[idx_i, idx_j] = np.exp(-(np.power(distance, 2)) / (2 * np.power(sigma, 2)))
201
202            # Constant
203            else:
204                kernel[idx_i, idx_j] = 1.0
205
206    return kernel
207
208
209def kernel_base(
210    radius: float,
211    circular: bool = False,
212    distance_weighted: bool = False,
213    normalised: bool = True,
214    hole: bool = False,
215    method: int = 0,
216    decay: float = 0.2,
217    sigma: float = 2.0,
218) -> np.ndarray:
219    """Creates a 2D kernel for convolution.
220
221    Parameters
222    ----------
223    radius : float
224        Radius of the kernel.
225
226    circular : bool
227        Whether to use a circular kernel.
228
229    distance_weighted : bool
230        Whether to use a distance weighted kernel.
231
232    normalised : bool
233        Whether to normalise the kernel.
234
235    hole : bool
236        Whether to create a hole in the center of the kernel.
237
238    method : int
239        Method to use for weighting.
240        0. linear
241        1. sqrt
242        2. power
243        3. gaussian
244        4. constant
245
246    decay : float
247        Decay rate for distance weighted kernels. Only used if `distance_weighted` is True.
248
249    sigma : float
250        Sigma for gaussian distance weighted kernels. Only used if `distance_weighted` is True and `method` is 3.
251
252    Returns
253    -------
254    kernel : np.ndarray
255        The kernel.
256    """
257    size = np.int64(np.ceil(radius) * 2 + 1)
258    kernel = np.ones((size, size), dtype=np.float32)
259
260    if hole:
261        kernel[size // 2, size // 2] = 0.0
262
263    if circular:
264        circular_kernel = _circular_kernel_2D(radius)
265        kernel *= circular_kernel
266
267    if distance_weighted:
268        distance_weighted_kernel = _distance_weighted_kernel_2D(
269            radius,
270            method,
271            decay,
272            sigma,
273        )
274        kernel *= distance_weighted_kernel
275
276    if normalised:
277        kernel /= np.sum(kernel)
278
279    return kernel
280
281
282
283@jit(nopython=True, nogil=True, cache=True, fastmath=True, inline='always')
284def kernel_shift(
285    x_offset: float,
286    y_offset: float,
287) -> Tuple[np.ndarray, np.ndarray]:
288    """Create a 2D shift kernel.
289
290    This function returns a kernel that can be used to shift a raster by a fractional
291    number of pixels in the x and y directions. The kernel can also be used to simulate
292    channel misalignment in image augmentation.
293
294    Parameters
295    ----------
296    x_offset : float
297        The horizontal (x) offset to apply.
298
299    y_offset : float
300        The vertical (y) offset to apply.
301
302    Returns
303    -------
304    Tuple[np.ndarray, np.ndarray]
305        A tuple of two numpy arrays. The first array contains the (x, y) offsets
306        of the kernel values. The second array contains the corresponding weights
307        of each kernel value.
308    """
309    if x_offset == 0.0 and y_offset == 0.0:
310        offsets = np.array([[0, 0]], dtype=np.int64)
311        weights = np.array([1.0], dtype=np.float32)
312
313        return offsets, weights
314
315    y0 = [int(np.floor(y_offset)), int(np.ceil(y_offset))] if y_offset != 0 else [0, 0]
316    x0 = [int(np.floor(x_offset)), int(np.ceil(x_offset))] if x_offset != 0 else [0, 0]
317
318    if x_offset == 0.0 or x_offset % 1 == 0.0:
319        offsets = np.zeros((2, 2), dtype=np.int64)
320        weights = np.zeros(2, dtype=np.float32)
321
322        offsets[0] = [int(x_offset) if x_offset % 1 == 0.0 else 0, y0[0]]
323        offsets[1] = [int(x_offset) if x_offset % 1 == 0.0 else 0, y0[1]]
324
325        weights[0] = y_offset - y0[0]
326        weights[1] = 1 - weights[0]
327
328    elif y_offset == 0.0 or y_offset % 1 == 0.0:
329        offsets = np.zeros((2, 2), dtype=np.int64)
330        weights = np.zeros(2, dtype=np.float32)
331
332        offsets[0] = [x0[0], int(y_offset) if y_offset % 1 == 0.0 else 0]
333        offsets[1] = [x0[1], int(y_offset) if y_offset % 1 == 0.0 else 0]
334
335        weights[0] = x_offset - x0[0]
336        weights[1] = 1 - weights[0]
337
338    else:
339        offsets = np.zeros((4, 2), dtype=np.int64)
340        weights = np.zeros(4, dtype=np.float32)
341
342        offsets[0] = [x0[0], y0[0]]
343        offsets[1] = [x0[0], y0[1]]
344        offsets[2] = [x0[1], y0[0]]
345        offsets[3] = [x0[1], y0[1]]
346
347        weights[0] = (1 - (x_offset - offsets[0][0])) * (1 - (y_offset - offsets[0][1]))
348        weights[1] = (1 - (x_offset - offsets[1][0])) * (1 + (y_offset - offsets[1][1]))
349        weights[2] = (1 + (x_offset - offsets[2][0])) * (1 - (y_offset - offsets[2][1]))
350        weights[3] = (1 + (x_offset - offsets[3][0])) * (1 + (y_offset - offsets[3][1]))
351
352    return offsets, weights
353
354
355@jit(nopython=True, nogil=True, cache=True, fastmath=True)
356def kernel_unsharp(
357    radius: float = 1.0,
358    intensity: float = 1.0,
359) -> np.ndarray:
360    """Create a 2D unsharp kernel.
361
362    This function returns a kernel that can be used to apply an unsharp mask to an image.
363
364    Parameters
365    ----------
366    radius : float, optional
367        The radius of the kernel. Default: 1.0.
368
369    intensity : float, optional
370        The intensity of the unsharp mask. Default: 1.0.
371
372    Returns
373    -------
374    np.ndarray
375        The kernel.
376    """
377    kernel = kernel_base(
378        radius=radius,
379        circular=True,
380        distance_weighted=True,
381        method=3,
382        normalised=True,
383        hole=True,
384    )
385
386    kernel *= intensity
387    kernel_sum = np.sum(kernel) + 1.0
388    kernel *= -1.0
389    kernel[kernel.shape[0] // 2, kernel.shape[1] // 2] = kernel_sum
390
391    return kernel
392
393
394@jit(nopython=True, nogil=True, cache=True, fastmath=True, parallel=True)
395def kernel_sobel(
396    radius=1,
397    scale=2,
398) -> Tuple[np.ndarray, np.ndarray]:
399    """Creates a 2D Sobel style kernel consisting of a horizontal and vertical component.
400    This function returns a kernel that can be used to apply a Sobel filter to an image.
401
402    `kernel_gx, kernel_gy = get_kernel_sobel(radius=1, scale=2)`
403
404    The kernels for radis=2, scale=2 are:
405    ```python
406    gx = [
407        [ 0.56  0.85  0.   -0.85 -0.56],
408        [ 0.85  1.5   0.   -1.5  -0.85],
409        [ 1.    2.    0.   -2.   -1.  ],
410        [ 0.85  1.5   0.   -1.5  -0.85],
411        [ 0.56  0.85  0.   -0.85 -0.56],
412    ]
413
414    gy = [
415        [ 0.56  0.85  1.    0.85  0.56],
416        [ 0.85  1.5   2.    1.5   0.85],
417        [ 0.    0.    0.    0.    0.  ],
418        [-0.85 -1.5  -2.   -1.5  -0.85],
419        [-0.56 -0.85 -1.   -0.85 -0.56],
420    ]
421    ```
422
423    Parameters
424    ----------
425    radius : float, optional
426        The radius of the kernel. Default: 1.0.
427
428    scale : float, optional
429        The scale of the kernel. Default: 2.0.
430
431    Returns
432    -------
433    Tuple[np.ndarray, np.ndarray]
434        The kernels
435
436    """
437    size = np.int64(np.ceil(radius) * 2 + 1)
438    kernel = np.zeros((size, size), dtype=np.float32)
439    center = np.array([0.0, 0.0], dtype=np.float32)
440
441    step = size // 2
442    for idx_i, col in enumerate(range(-step, step + 1)):
443        for idx_j, row in enumerate(range(-step, step + 1)):
444            point = np.array([col, row], dtype=np.float32)
445            distance = _distance_2D(center, point)
446            if col == 0 and row == 0:
447                kernel[idx_i, idx_j] = 0
448            else:
449                weight = np.power((1 - 0.5), distance) * 2
450                kernel[idx_i, idx_j] = weight * scale
451
452    # vertical
453    kernel_gx = kernel.copy()
454    kernel_gx[:, size // 2:] *= -1
455    kernel_gx[:, size // 2] = 0
456
457    # horisontal
458    kernel_gy = kernel.copy()
459    kernel_gy[size // 2:, :] *= -1
460    kernel_gy[size // 2, :] = 0
461
462    return kernel_gx, kernel_gy
463
464
465@jit(nopython=True, nogil=True, cache=True, fastmath=True)
466def kernel_get_offsets_and_weights(
467    kernel: np.ndarray,
468    remove_zero_weights: bool = True,
469) -> Tuple[np.ndarray, np.ndarray, int]:
470    """Generates a list of offsets and weights for a given kernel.
471
472    Parameters
473    ----------
474    kernel : np.ndarray
475        The kernel to generate offsets and weights for.
476
477    remove_zero_weights : bool, optional
478        Whether to remove offsets and weights with zero weights. Default: True.
479
480    Returns
481    -------
482    Tuple[np.ndarray, np.ndarray]
483        The offsets and weights
484    """
485    count = np.count_nonzero(kernel) if remove_zero_weights else kernel.size
486
487    offsets = np.zeros((count, 2), dtype=np.int64)
488    weights = np.zeros((count), dtype=np.float32)
489
490    step = kernel.shape[0] // 2
491    index = 0
492    for col in range(-step, step + 1):
493        for row in range(-step, step + 1):
494            if kernel[col + step, row + step] != 0.0:
495                offsets[index][0] = col
496                offsets[index][1] = row
497                weights[index] = kernel[col + step, row + step]
498
499                index += 1
500
501    return offsets, weights