vkit.mechanism.distortion.photometric.noise

  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 Any, Optional, Mapping
 15
 16import attrs
 17import numpy as np
 18from numpy.random import Generator as RandomGenerator
 19
 20from vkit.element import Image
 21from ..interface import DistortionConfig, DistortionNopState, Distortion
 22from .opt import extract_mat_from_image, clip_mat_back_to_uint8
 23
 24
 25@attrs.define
 26class GaussionNoiseConfig(DistortionConfig):
 27    std: float
 28
 29    _rng_state: Optional[Mapping[str, Any]] = None
 30
 31    @property
 32    def supports_rng_state(self) -> bool:
 33        return True
 34
 35    @property
 36    def rng_state(self) -> Optional[Mapping[str, Any]]:
 37        return self._rng_state
 38
 39    @rng_state.setter
 40    def rng_state(self, val: Mapping[str, Any]):
 41        self._rng_state = val
 42
 43
 44def gaussion_noise_image(
 45    config: GaussionNoiseConfig,
 46    state: Optional[DistortionNopState[GaussionNoiseConfig]],
 47    image: Image,
 48    rng: Optional[RandomGenerator],
 49):
 50    assert rng
 51    mat = extract_mat_from_image(image, np.int16)
 52    noise = np.round(rng.normal(0, config.std, mat.shape)).astype(np.int16)
 53    mat = clip_mat_back_to_uint8(mat + noise)
 54    return Image(mat=mat)
 55
 56
 57gaussion_noise = Distortion(
 58    config_cls=GaussionNoiseConfig,
 59    state_cls=DistortionNopState[GaussionNoiseConfig],
 60    func_image=gaussion_noise_image,
 61)
 62
 63
 64@attrs.define
 65class PoissonNoiseConfig(DistortionConfig):
 66    _rng_state: Optional[Mapping[str, Any]] = None
 67
 68    @property
 69    def supports_rng_state(self) -> bool:
 70        return True
 71
 72    @property
 73    def rng_state(self) -> Optional[Mapping[str, Any]]:
 74        return self._rng_state
 75
 76    @rng_state.setter
 77    def rng_state(self, val: Mapping[str, Any]):
 78        self._rng_state = val
 79
 80
 81def poisson_noise_image(
 82    config: PoissonNoiseConfig,
 83    state: Optional[DistortionNopState[PoissonNoiseConfig]],
 84    image: Image,
 85    rng: Optional[RandomGenerator],
 86):
 87    assert rng
 88    mat = rng.poisson(extract_mat_from_image(image, np.float32))
 89    mat = clip_mat_back_to_uint8(mat)
 90    return Image(mat=mat)
 91
 92
 93poisson_noise = Distortion(
 94    config_cls=PoissonNoiseConfig,
 95    state_cls=DistortionNopState[PoissonNoiseConfig],
 96    func_image=poisson_noise_image,
 97)
 98
 99
100@attrs.define
101class ImpulseNoiseConfig(DistortionConfig):
102    prob_salt: float
103    prob_pepper: float
104
105    _rng_state: Optional[Mapping[str, Any]] = None
106
107    @property
108    def supports_rng_state(self) -> bool:
109        return True
110
111    @property
112    def rng_state(self) -> Optional[Mapping[str, Any]]:
113        return self._rng_state
114
115    @rng_state.setter
116    def rng_state(self, val: Mapping[str, Any]):
117        self._rng_state = val
118
119
120def impulse_noise_image(
121    config: ImpulseNoiseConfig,
122    state: Optional[DistortionNopState[ImpulseNoiseConfig]],
123    image: Image,
124    rng: Optional[RandomGenerator],
125):
126    assert rng
127
128    # https://www.programmersought.com/article/3363136769/
129    prob_presv = 1 - config.prob_salt - config.prob_pepper
130    assert prob_presv >= 0.0
131
132    mask = rng.choice(
133        (0, 1, 2),
134        size=image.shape,
135        p=[prob_presv, config.prob_salt, config.prob_pepper],
136    )
137
138    mat = image.mat.copy()
139    # Salt.
140    mat[mask == 1] = 255
141    # Pepper.
142    mat[mask == 2] = 0
143
144    return Image(mat=mat)
145
146
147impulse_noise = Distortion(
148    config_cls=ImpulseNoiseConfig,
149    state_cls=DistortionNopState[ImpulseNoiseConfig],
150    func_image=impulse_noise_image,
151)
152
153
154@attrs.define
155class SpeckleNoiseConfig(DistortionConfig):
156    std: float
157
158    _rng_state: Optional[Mapping[str, Any]] = None
159
160    @property
161    def supports_rng_state(self) -> bool:
162        return True
163
164    @property
165    def rng_state(self) -> Optional[Mapping[str, Any]]:
166        return self._rng_state
167
168    @rng_state.setter
169    def rng_state(self, val: Mapping[str, Any]):
170        self._rng_state = val
171
172
173def speckle_noise_image(
174    config: SpeckleNoiseConfig,
175    state: Optional[DistortionNopState[SpeckleNoiseConfig]],
176    image: Image,
177    rng: Optional[RandomGenerator],
178):
179    assert rng
180    mat = extract_mat_from_image(image, np.float32)
181    noise = rng.normal(0, config.std, mat.shape)
182    mat = clip_mat_back_to_uint8(mat + mat * noise)
183    return Image(mat=mat)
184
185
186speckle_noise = Distortion(
187    config_cls=SpeckleNoiseConfig,
188    state_cls=DistortionNopState[SpeckleNoiseConfig],
189    func_image=speckle_noise_image,
190)