vkit.mechanism.distortion.photometric.blur

  1# Copyright 2022 vkit-x Administrator. All Rights Reserved.
  2#
  3# This project (vkit-x/vkit) is dual-licensed under commercial and SSPL licenses.
  4#
  5# The commercial license gives you the full rights to create and distribute software
  6# on your own terms without any SSPL license obligations. For more information,
  7# please see the "LICENSE_COMMERCIAL.txt" file.
  8#
  9# This project is also available under Server Side Public License (SSPL).
 10# The SSPL licensing is ideal for use cases such as open source projects with
 11# SSPL distribution, student/academic purposes, hobby projects, internal research
 12# projects without external distribution, or other projects where all SSPL
 13# obligations can be met. For more information, please see the "LICENSE_SSPL.txt" file.
 14from typing import Optional, Tuple, Mapping, Any
 15
 16import attrs
 17import numpy as np
 18from numpy.random import Generator as RandomGenerator
 19import cv2 as cv
 20
 21from vkit.element import Image
 22from ..interface import DistortionConfig, DistortionNopState, Distortion
 23from .opt import to_rgb_image, to_original_image, clip_mat_back_to_uint8
 24
 25
 26def _estimate_gaussian_kernel_size(sigma: float):
 27    kernel_size = max(3, round(3 * sigma) + 1)
 28    if kernel_size % 2 == 0:
 29        kernel_size += 1
 30    return kernel_size
 31
 32
 33def _get_anti_aliasing_kernel_size_and_padding(anti_aliasing_sigma: float):
 34    kernel_size = _estimate_gaussian_kernel_size(anti_aliasing_sigma)
 35
 36    anti_aliasing_ksize = (kernel_size, kernel_size)
 37    anti_aliasing_kernel_padding = kernel_size // 2 * 2
 38    return anti_aliasing_ksize, anti_aliasing_kernel_padding
 39
 40
 41def _apply_anti_aliasing_to_kernel(
 42    kernel: np.ndarray,
 43    anti_aliasing_ksize: Tuple[int, int],
 44    anti_aliasing_sigma: float,
 45):
 46    return cv.GaussianBlur(kernel, anti_aliasing_ksize, anti_aliasing_sigma)
 47
 48
 49@attrs.define
 50class GaussianBlurConfig(DistortionConfig):
 51    sigma: float
 52
 53
 54def gaussian_blur_image(
 55    config: GaussianBlurConfig,
 56    state: Optional[DistortionNopState[GaussianBlurConfig]],
 57    image: Image,
 58    rng: Optional[RandomGenerator],
 59):
 60    mode = image.mode
 61    image = to_rgb_image(image, mode)
 62
 63    kernel_size = _estimate_gaussian_kernel_size(config.sigma)
 64    ksize = (kernel_size, kernel_size)
 65    mat = cv.GaussianBlur(image.mat, ksize, config.sigma)
 66    image = attrs.evolve(image, mat=mat)
 67
 68    image = to_original_image(image, mode)
 69    return image
 70
 71
 72gaussian_blur = Distortion(
 73    config_cls=GaussianBlurConfig,
 74    state_cls=DistortionNopState[GaussianBlurConfig],
 75    func_image=gaussian_blur_image,
 76)
 77
 78
 79@attrs.define
 80class DefocusBlurConfig(DistortionConfig):
 81    radius: int
 82    anti_aliasing_sigma: float = 0.5
 83
 84
 85def defocus_blur_image(
 86    config: DefocusBlurConfig,
 87    state: Optional[DistortionNopState[DefocusBlurConfig]],
 88    image: Image,
 89    rng: Optional[RandomGenerator],
 90):
 91    # Generate blurring kernel.
 92    assert 0 < config.radius
 93    kernel_size = 2 * config.radius + 1
 94
 95    (
 96        anti_aliasing_ksize,
 97        anti_aliasing_kernel_padding,
 98    ) = _get_anti_aliasing_kernel_size_and_padding(config.anti_aliasing_sigma)
 99    kernel_size += anti_aliasing_kernel_padding
100
101    begin = -(kernel_size // 2)
102    end = begin + kernel_size
103    coords = np.arange(begin, end)
104    x, y = np.meshgrid(coords, coords)
105    kernel: np.ndarray = ((x**2 + y**2) <= config.radius**2).astype(np.float32)
106    kernel /= kernel.sum()
107
108    kernel = _apply_anti_aliasing_to_kernel(
109        kernel,
110        anti_aliasing_ksize,
111        config.anti_aliasing_sigma,
112    )
113
114    # Convolution.
115    mode = image.mode
116    image = to_rgb_image(image, mode)
117
118    mat = cv.filter2D(image.mat, -1, kernel)
119    image = attrs.evolve(image, mat=mat)
120
121    image = to_original_image(image, mode)
122    return image
123
124
125defocus_blur = Distortion(
126    config_cls=DefocusBlurConfig,
127    state_cls=DistortionNopState[DefocusBlurConfig],
128    func_image=defocus_blur_image,
129)
130
131
132@attrs.define
133class MotionBlurConfig(DistortionConfig):
134    radius: int
135    angle: int
136    anti_aliasing_sigma: float = 0.5
137
138
139def motion_blur_image(
140    config: MotionBlurConfig,
141    state: Optional[DistortionNopState[MotionBlurConfig]],
142    image: Image,
143    rng: Optional[RandomGenerator],
144):
145    # Generate blurring kernel.
146    kernel_size = 2 * config.radius + 1
147
148    (
149        anti_aliasing_ksize,
150        anti_aliasing_kernel_padding,
151    ) = _get_anti_aliasing_kernel_size_and_padding(config.anti_aliasing_sigma)
152    anti_aliasing_kernel_padding_half = anti_aliasing_kernel_padding // 2
153    center = config.radius + anti_aliasing_kernel_padding_half
154    left = anti_aliasing_kernel_padding_half
155    right = left + kernel_size - 1
156    kernel_size += anti_aliasing_kernel_padding
157
158    kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)
159    kernel[center, left:right + 1] = 1.0
160
161    trans_mat = cv.getRotationMatrix2D(
162        (center, center),
163        # 1. to [0, 359].
164        # 2. getRotationMatrix2D accepts counter-clockwise angle, hence need to subtract.
165        360 - (config.angle % 360),
166        1.0,
167    )
168    kernel = cv.warpAffine(kernel, trans_mat, kernel.shape)
169    kernel /= kernel.sum()
170
171    kernel = _apply_anti_aliasing_to_kernel(
172        kernel,
173        anti_aliasing_ksize,
174        config.anti_aliasing_sigma,
175    )
176
177    # Convolution.
178    mode = image.mode
179    image = to_rgb_image(image, mode)
180
181    mat = cv.filter2D(image.mat, -1, kernel)
182    image = attrs.evolve(image, mat=mat)
183
184    image = to_original_image(image, mode)
185    return image
186
187
188motion_blur = Distortion(
189    config_cls=MotionBlurConfig,
190    state_cls=DistortionNopState[MotionBlurConfig],
191    func_image=motion_blur_image,
192)
193
194
195@attrs.define
196class GlassBlurConfig(DistortionConfig):
197    sigma: float
198    delta: int = 1
199    loop: int = 5
200
201    _rng_state: Optional[Mapping[str, Any]] = None
202
203    @property
204    def supports_rng_state(self) -> bool:
205        return True
206
207    @property
208    def rng_state(self) -> Optional[Mapping[str, Any]]:
209        return self._rng_state
210
211    @rng_state.setter
212    def rng_state(self, val: Mapping[str, Any]):
213        self._rng_state = val
214
215
216def glass_blur_image(
217    config: GlassBlurConfig,
218    state: Optional[DistortionNopState[GlassBlurConfig]],
219    image: Image,
220    rng: Optional[RandomGenerator],
221):
222    mode = image.mode
223    image = to_rgb_image(image, mode)
224
225    # Gaussian blur.
226    kernel_size = _estimate_gaussian_kernel_size(config.sigma)
227    ksize = (kernel_size, kernel_size)
228    mat = cv.GaussianBlur(image.mat, ksize, config.sigma)
229
230    # Random pixel swap.
231    assert rng is not None
232
233    coords_y = np.arange(image.height)
234    coords_x = np.arange(image.width)
235    pos_x, pos_y = np.meshgrid(coords_x, coords_y)
236
237    for _ in range(config.loop):
238        offset_y = rng.integers(0, 2 * config.delta + 1)
239        center_y = np.arange(offset_y, image.height - config.delta, 2 * config.delta + 1)
240        num_center_y = center_y.shape[0]
241        center_y = center_y.reshape(-1, 1)
242
243        offset_x = rng.integers(0, 2 * config.delta + 1)
244        center_x = np.arange(offset_x, image.width - config.delta, 2 * config.delta + 1)
245        num_center_x = center_x.shape[0]
246        center_x = center_x.reshape(1, -1)
247
248        delta_shape = (num_center_y, num_center_x)
249        delta_y = rng.integers(-config.delta, config.delta + 1, delta_shape)
250        delta_x = rng.integers(-config.delta, config.delta + 1, delta_shape)
251
252        deformed_y: np.ndarray = pos_y[center_y, center_x] + delta_y
253        deformed_y = np.clip(deformed_y, 0, image.height - 1)
254        deformed_x: np.ndarray = pos_x[center_y, center_x] + delta_x
255        deformed_x = np.clip(deformed_x, 0, image.width - 1)
256
257        # Swap.
258        pos_y[center_y, center_x], pos_y[deformed_y, deformed_x] = \
259            pos_y[deformed_y, deformed_x], pos_y[center_y, center_x]
260
261        pos_x[center_y, center_x], pos_x[deformed_y, deformed_x] = \
262            pos_x[deformed_y, deformed_x], pos_x[center_y, center_x]
263
264    mat = mat[pos_y, pos_x]
265    image = attrs.evolve(image, mat=mat)
266
267    image = to_original_image(image, mode)
268    return image
269
270
271glass_blur = Distortion(
272    config_cls=GlassBlurConfig,
273    state_cls=DistortionNopState[GlassBlurConfig],
274    func_image=glass_blur_image,
275)
276
277
278@attrs.define
279class ZoomInBlurConfig(DistortionConfig):
280    ratio: float = 0.1
281    step: float = 0.01
282    alpha: float = 0.5
283
284
285def zoom_in_blur_image(
286    config: ZoomInBlurConfig,
287    state: Optional[DistortionNopState[ZoomInBlurConfig]],
288    image: Image,
289    rng: Optional[RandomGenerator],
290):
291    mode = image.mode
292    image = to_rgb_image(image, mode)
293
294    mat = image.mat.astype(np.uint16)
295    count = 1
296    for ratio in np.arange(
297        1 + config.step,
298        1 + config.ratio + config.step,
299        config.step,
300    ):
301        resized_height = round(image.height * ratio)
302        resized_width = round(image.width * ratio)
303        resized_image = image.to_resized_image(resized_height, resized_width)
304
305        pad_y = resized_height - image.height
306        up = pad_y // 2
307        down = up + image.height - 1
308
309        pad_x = resized_width - image.width
310        left = pad_x // 2
311        right = left + image.width - 1
312
313        mat += resized_image.mat[up:down + 1, left:right + 1]
314        count += 1
315
316    # NOTE: Label confusion could be significant if alpha is large.
317    mat: np.ndarray = (1 - config.alpha) * image.mat + config.alpha * np.round(mat / count)
318    mat = clip_mat_back_to_uint8(mat)
319
320    image = attrs.evolve(image, mat=mat)
321
322    image = to_original_image(image, mode)
323    return image
324
325
326zoom_in_blur = Distortion(
327    config_cls=ZoomInBlurConfig,
328    state_cls=DistortionNopState[ZoomInBlurConfig],
329    func_image=zoom_in_blur_image,
330)