vkit.mechanism.distortion_policy.geometric.affine

  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 Tuple
 15
 16import attrs
 17from numpy.random import Generator as RandomGenerator
 18
 19from vkit.mechanism import distortion
 20from ..type import DistortionConfigGenerator, DistortionPolicyFactory
 21from ..opt import sample_int, sample_float
 22
 23
 24@attrs.define
 25class ShearHoriConfigGeneratorConfig:
 26    angle_min: int = 1
 27    angle_max: int = 30
 28    prob_negative: float = 0.5
 29
 30
 31class ShearHoriConfigGenerator(
 32    DistortionConfigGenerator[
 33        ShearHoriConfigGeneratorConfig,
 34        distortion.ShearHoriConfig,
 35    ]
 36):  # yapf: disable
 37
 38    def __call__(self, shape: Tuple[int, int], rng: RandomGenerator):
 39        angle = sample_int(
 40            level=self.level,
 41            value_min=self.config.angle_min,
 42            value_max=self.config.angle_max,
 43            prob_negative=self.config.prob_negative,
 44            rng=rng,
 45        )
 46
 47        return distortion.ShearHoriConfig(angle=angle)
 48
 49
 50shear_hori_policy_factory = DistortionPolicyFactory(
 51    distortion.shear_hori,
 52    ShearHoriConfigGenerator,
 53)
 54
 55
 56@attrs.define
 57class ShearVertConfigGeneratorConfig:
 58    angle_min: int = 1
 59    angle_max: int = 30
 60    prob_negative: float = 0.5
 61
 62
 63class ShearVertConfigGenerator(
 64    DistortionConfigGenerator[
 65        ShearVertConfigGeneratorConfig,
 66        distortion.ShearVertConfig,
 67    ]
 68):  # yapf: disable
 69
 70    def __call__(self, shape: Tuple[int, int], rng: RandomGenerator):
 71        angle = sample_int(
 72            level=self.level,
 73            value_min=self.config.angle_min,
 74            value_max=self.config.angle_max,
 75            prob_negative=self.config.prob_negative,
 76            rng=rng,
 77        )
 78
 79        return distortion.ShearVertConfig(angle=angle)
 80
 81
 82shear_vert_policy_factory = DistortionPolicyFactory(
 83    distortion.shear_vert,
 84    ShearVertConfigGenerator,
 85)
 86
 87
 88@attrs.define
 89class RotateConfigGeneratorConfig:
 90    angle_min: int = 1
 91    angle_max: int = 180
 92    prob_negative: float = 0.5
 93
 94
 95class RotateConfigGenerator(
 96    DistortionConfigGenerator[
 97        RotateConfigGeneratorConfig,
 98        distortion.RotateConfig,
 99    ]
100):  # yapf: disable
101
102    def __call__(self, shape: Tuple[int, int], rng: RandomGenerator):
103        angle = sample_int(
104            level=self.level,
105            value_min=self.config.angle_min,
106            value_max=self.config.angle_max,
107            prob_negative=self.config.prob_negative,
108            rng=rng,
109        )
110        return distortion.RotateConfig(angle=angle)
111
112
113rotate_policy_factory = DistortionPolicyFactory(
114    distortion.rotate,
115    RotateConfigGenerator,
116)
117
118
119@attrs.define
120class SkewHoriConfigGeneratorConfig:
121    ratio_min: float = 0.0
122    ratio_max: float = 0.35
123    prob_negative: float = 0.5
124
125
126class SkewHoriConfigGenerator(
127    DistortionConfigGenerator[
128        SkewHoriConfigGeneratorConfig,
129        distortion.SkewHoriConfig,
130    ]
131):  # yapf: disable
132
133    def __call__(self, shape: Tuple[int, int], rng: RandomGenerator):
134        ratio = sample_float(
135            level=self.level,
136            value_min=self.config.ratio_min,
137            value_max=self.config.ratio_max,
138            prob_reciprocal=None,
139            rng=rng,
140        )
141        if rng.random() < self.config.prob_negative:
142            ratio *= -1
143
144        return distortion.SkewHoriConfig(ratio=ratio)
145
146
147skew_hori_policy_factory = DistortionPolicyFactory(
148    distortion.skew_hori,
149    SkewHoriConfigGenerator,
150)
151
152
153@attrs.define
154class SkewVertConfigGeneratorConfig:
155    ratio_min: float = 0.0
156    ratio_max: float = 0.35
157    prob_negative: float = 0.5
158
159
160class SkewVertConfigGenerator(
161    DistortionConfigGenerator[
162        SkewVertConfigGeneratorConfig,
163        distortion.SkewVertConfig,
164    ]
165):  # yapf: disable
166
167    def __call__(self, shape: Tuple[int, int], rng: RandomGenerator):
168        ratio = sample_float(
169            level=self.level,
170            value_min=self.config.ratio_min,
171            value_max=self.config.ratio_max,
172            prob_reciprocal=None,
173            rng=rng,
174        )
175        if rng.random() < self.config.prob_negative:
176            ratio *= -1
177
178        return distortion.SkewVertConfig(ratio=ratio)
179
180
181skew_vert_policy_factory = DistortionPolicyFactory(
182    distortion.skew_vert,
183    SkewVertConfigGenerator,
184)