vkit.mechanism.distortion.photometric.effect

  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, ImageMode
 22from ..interface import DistortionConfig, DistortionNopState, Distortion
 23from .opt import to_rgb_image, to_original_image, clip_mat_back_to_uint8
 24
 25
 26@attrs.define
 27class JpegQualityConfig(DistortionConfig):
 28    quality: int
 29
 30
 31def jpeg_quality_image(
 32    config: JpegQualityConfig,
 33    state: Optional[DistortionNopState[JpegQualityConfig]],
 34    image: Image,
 35    rng: Optional[RandomGenerator],
 36):
 37    mode = image.mode
 38    image = to_rgb_image(image, mode)
 39
 40    assert 0 <= config.quality <= 100
 41    _, buffer = cv.imencode('.jpeg', image.mat, [cv.IMWRITE_JPEG_QUALITY, config.quality])
 42    mat: np.ndarray = cv.imdecode(buffer, cv.IMREAD_UNCHANGED)
 43    image = attrs.evolve(image, mat=mat)
 44
 45    image = to_original_image(image, mode)
 46    return image
 47
 48
 49jpeg_quality = Distortion(
 50    config_cls=JpegQualityConfig,
 51    state_cls=DistortionNopState[JpegQualityConfig],
 52    func_image=jpeg_quality_image,
 53)
 54
 55
 56@attrs.define
 57class PixelationConfig(DistortionConfig):
 58    ratio: float
 59
 60
 61def pixelation_image(
 62    config: PixelationConfig,
 63    state: Optional[DistortionNopState[PixelationConfig]],
 64    image: Image,
 65    rng: Optional[RandomGenerator],
 66):
 67    # Downsample.
 68    assert 0 < config.ratio < 1
 69    resized_height = round(image.height * config.ratio)
 70    resized_width = round(image.width * config.ratio)
 71    dsize = (resized_width, resized_height)
 72    mat: np.ndarray = cv.resize(image.mat, dsize, interpolation=cv.INTER_LINEAR)
 73
 74    # Upsample.
 75    dsize = (image.width, image.height)
 76    mat: np.ndarray = cv.resize(mat, dsize, interpolation=cv.INTER_NEAREST)
 77
 78    image = attrs.evolve(image, mat=mat)
 79    return image
 80
 81
 82pixelation = Distortion(
 83    config_cls=PixelationConfig,
 84    state_cls=DistortionNopState[PixelationConfig],
 85    func_image=pixelation_image,
 86)
 87
 88
 89def generate_diamond_square_mask(
 90    shape: Tuple[int, int],
 91    roughness: float,
 92    rng: RandomGenerator,
 93):
 94    assert 0.0 <= roughness <= 1.0
 95
 96    height, width = shape
 97    size = int(2**np.ceil(np.log2(max(height, width))) + 1)
 98
 99    mask = np.zeros((size, size), dtype=np.float32)
100    # Initialize 4 corners.
101    mask[0, 0] = rng.uniform(0.0, 1.0)
102    mask[0, -1] = rng.uniform(0.0, 1.0)
103    mask[-1, -1] = rng.uniform(0.0, 1.0)
104    mask[-1, 0] = rng.uniform(0.0, 1.0)
105
106    step = size - 1
107    iteration = 0
108    while step >= 2:
109        step_roughness = roughness**iteration
110
111        squares: np.ndarray = mask[0:size:step, 0:size:step]
112
113        square_sum_vert = squares + np.roll(squares, shift=-1, axis=0)
114        square_sum_hori = squares + np.roll(squares, shift=-1, axis=1)
115
116        # Diamond step.
117        square_sum = square_sum_vert + square_sum_hori
118        square_sum = square_sum[:-1, :-1]
119        diamonds = ((1 - step_roughness) * square_sum / 4
120                    + step_roughness * rng.uniform(0, 1, square_sum.shape))
121        mask[step // 2:size:step, step // 2:size:step] = diamonds
122
123        # Square step.
124        diamond_sum_vert = diamonds + np.roll(diamonds, shift=1, axis=0)
125        diamond_sum_vert = np.vstack([diamond_sum_vert, diamond_sum_vert[0]])
126        square_sum0 = square_sum_hori[:, :-1] + diamond_sum_vert
127        squares0 = ((1 - step_roughness) * square_sum0 / 4
128                    + step_roughness * rng.uniform(0, 1, square_sum0.shape))
129        mask[0:size:step, step // 2:size:step] = squares0
130
131        diamond_sum_hori = diamonds + np.roll(diamonds, shift=1, axis=1)
132        diamond_sum_hori = np.hstack([diamond_sum_hori, diamond_sum_hori[0].reshape(-1, 1)])
133        square_sum1 = square_sum_vert[:-1] + diamond_sum_hori
134        squares1 = ((1 - step_roughness) * square_sum1 / 4
135                    + step_roughness * rng.uniform(0, 1, square_sum1.shape))
136        mask[step // 2:size:step, 0:size:step] = squares1
137
138        iteration += 1
139        step = step // 2
140
141    up = rng.integers(0, size - height + 1)
142    down = up + height - 1
143    left = rng.integers(0, size - width + 1)
144    right = left + width - 1
145    return mask[up:down + 1, left:right + 1]
146
147
148@attrs.define
149class FogConfig(DistortionConfig):
150    roughness: float
151    fog_rgb: Tuple[int, int, int] = (226, 238, 234)
152    ratio_max: float = 1.0
153    ratio_min: float = 0.0
154
155    _rng_state: Optional[Mapping[str, Any]] = None
156
157    @property
158    def supports_rng_state(self) -> bool:
159        return True
160
161    @property
162    def rng_state(self) -> Optional[Mapping[str, Any]]:
163        return self._rng_state
164
165    @rng_state.setter
166    def rng_state(self, val: Mapping[str, Any]):
167        self._rng_state = val
168
169
170def fog_image(
171    config: FogConfig,
172    state: Optional[DistortionNopState[FogConfig]],
173    image: Image,
174    rng: Optional[RandomGenerator],
175):
176    mode = image.mode
177    image = to_rgb_image(image, mode)
178
179    assert rng is not None
180    mask = generate_diamond_square_mask(
181        image.shape,
182        config.roughness,
183        rng,
184    )
185    # Equalize mask.
186    mask -= mask.min()
187    mask /= mask.max()
188    assert config.ratio_min < config.ratio_max
189    mask *= (config.ratio_max - config.ratio_min)
190    mask += config.ratio_min
191
192    mat = image.mat.astype(np.float32)
193
194    if image.mode == ImageMode.GRAYSCALE:
195        val = 0.2126 * config.fog_rgb[0] + 0.7152 * config.fog_rgb[1] + 0.0722 * config.fog_rgb[2]
196        fog = np.full(image.shape, val, dtype=np.float32)
197        mat = (1 - mask) * mat + mask * fog
198
199    else:
200        assert image.mode == ImageMode.RGB
201        fog = np.full((*image.shape, 3), config.fog_rgb, dtype=np.float32)
202        mask = np.expand_dims(mask, axis=-1)
203        mat = (1 - mask) * mat + mask * fog
204
205    mat = clip_mat_back_to_uint8(mat)
206    image = attrs.evolve(image, mat=mat)
207
208    image = to_original_image(image, mode)
209    return image
210
211
212fog = Distortion(
213    config_cls=FogConfig,
214    state_cls=DistortionNopState[FogConfig],
215    func_image=fog_image,
216)