vkit.mechanism.distortion.photometric.streak

  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, Optional, List
 15
 16import attrs
 17from numpy.random import Generator as RandomGenerator
 18import cv2 as cv
 19
 20from vkit.element import Image, Mask, Box
 21from ..interface import DistortionConfig, DistortionNopState, Distortion
 22
 23
 24def fill_vert_dash_gap(dash_thickness: int, dash_gap: int, mask: Mask):
 25    if dash_thickness <= 0 or dash_gap <= 0:
 26        return
 27
 28    with mask.writable_context:
 29        step = dash_thickness + dash_gap
 30        for offset_y in range(dash_gap):
 31            mask.mat[offset_y::step] = 0
 32
 33
 34def fill_hori_dash_gap(dash_thickness: int, dash_gap: int, mask: Mask):
 35    if dash_thickness <= 0 or dash_gap <= 0:
 36        return
 37
 38    with mask.writable_context:
 39        step = dash_thickness + dash_gap
 40        for offset_x in range(dash_gap):
 41            mask.mat[:, offset_x::step] = 0
 42
 43
 44@attrs.define
 45class LineStreakConfig(DistortionConfig):
 46    thickness: int = 1
 47    gap: int = 4
 48    dash_thickness: int = 0
 49    dash_gap: int = 0
 50    color: Tuple[int, int, int] = (0, 0, 0)
 51    alpha: float = 1.0
 52    enable_vert: bool = True
 53    enable_hori: bool = True
 54
 55
 56def line_streak_image(
 57    config: LineStreakConfig,
 58    state: Optional[DistortionNopState[LineStreakConfig]],
 59    image: Image,
 60    rng: Optional[RandomGenerator],
 61):
 62    masks: List[Mask] = []
 63
 64    step = config.thickness + config.gap
 65
 66    if config.enable_vert:
 67        mask = Mask.from_shapable(image)
 68
 69        with mask.writable_context:
 70            for offset_x in range(config.thickness):
 71                mask.mat[:, offset_x::step] = 1
 72
 73        fill_vert_dash_gap(
 74            dash_thickness=config.dash_thickness,
 75            dash_gap=config.dash_gap,
 76            mask=mask,
 77        )
 78
 79        masks.append(mask)
 80
 81    if config.enable_hori:
 82        mask = Mask.from_shapable(image)
 83
 84        with mask.writable_context:
 85            for offset_y in range(config.thickness):
 86                mask.mat[offset_y::step] = 1
 87
 88        fill_hori_dash_gap(
 89            dash_thickness=config.dash_thickness,
 90            dash_gap=config.dash_gap,
 91            mask=mask,
 92        )
 93
 94        masks.append(mask)
 95
 96    image = image.copy()
 97    for mask in masks:
 98        mask.fill_image(image, config.color, alpha=config.alpha)
 99    return image
100
101
102line_streak = Distortion(
103    config_cls=LineStreakConfig,
104    state_cls=DistortionNopState[LineStreakConfig],
105    func_image=line_streak_image,
106)
107
108
109def generate_centered_boxes(
110    height: int,
111    width: int,
112    aspect_ratio: float,
113    short_side_min: int,
114    short_side_step: int,
115):
116    center_y = height // 2
117    center_x = width // 2
118
119    boxes: List[Box] = []
120    idx = 0
121    while True:
122        short_side = short_side_min + idx * short_side_step
123        if aspect_ratio >= 1:
124            # hori side is longer.
125            height_min = short_side
126            width_min = round(height_min * aspect_ratio)
127        elif 0 < aspect_ratio < 1:
128            # vert side is longer.
129            width_min = short_side
130            height_min = round(width_min / aspect_ratio)
131        else:
132            raise NotImplementedError()
133
134        up = center_y - height_min // 2
135        down = up + height_min - 1
136        left = center_x - width_min // 2
137        right = left + width_min - 1
138
139        if (0 <= up and down < height) or (0 <= left and right < width):
140            boxes.append(Box(up=up, down=down, left=left, right=right))
141            idx += 1
142        else:
143            break
144
145    return boxes
146
147
148@attrs.define
149class RectangleStreakConfig(DistortionConfig):
150    thickness: int = 1
151    aspect_ratio: Optional[float] = None
152    dash_thickness: int = 0
153    dash_gap: int = 0
154    short_side_min: int = 10
155    short_side_step: int = 10
156    color: Tuple[int, int, int] = (0, 0, 0)
157    alpha: float = 1.0
158
159
160def rectangle_streak_image(
161    config: RectangleStreakConfig,
162    state: Optional[DistortionNopState[RectangleStreakConfig]],
163    image: Image,
164    rng: Optional[RandomGenerator],
165):
166    aspect_ratio = config.aspect_ratio
167    if aspect_ratio is None:
168        aspect_ratio = image.width / image.height
169
170    boxes = generate_centered_boxes(
171        height=image.height,
172        width=image.width,
173        aspect_ratio=aspect_ratio,
174        short_side_min=config.short_side_min,
175        short_side_step=config.short_side_step,
176    )
177
178    # Generate bars.
179    vert_bars: List[Box] = []
180    hori_bars: List[Box] = []
181
182    for box in boxes:
183        inner_up = box.down - config.thickness + 1
184        inner_down = box.up + config.thickness - 1
185        inner_left = box.right - config.thickness + 1
186        inner_right = box.left + config.thickness - 1
187
188        # Shared by left/right bars.
189        bar_up = max(0, box.up)
190        bar_down = min(image.height - 1, box.down)
191
192        # Left bar.
193        bar_left = max(0, box.left)
194        bar_right = inner_right
195
196        if 0 <= bar_right < image.width and bar_up <= bar_down:
197            vert_bars.append(Box(
198                up=bar_up,
199                down=bar_down,
200                left=bar_left,
201                right=bar_right,
202            ))
203
204        # Right bar.
205        bar_left = inner_left
206        bar_right = min(image.width - 1, box.right)
207
208        if 0 <= bar_left < image.width and bar_up <= bar_down:
209            vert_bars.append(Box(
210                up=bar_up,
211                down=bar_down,
212                left=bar_left,
213                right=bar_right,
214            ))
215
216        # Shared by top/bottom bars.
217        bar_left = max(0, inner_right + 1)
218        bar_right = min(image.width - 1, inner_left - 1)
219
220        # Top bar.
221        bar_up = max(0, box.up)
222        bar_down = inner_down
223
224        if 0 <= inner_down < image.height and bar_left <= bar_right:
225            hori_bars.append(Box(
226                up=bar_up,
227                down=bar_down,
228                left=bar_left,
229                right=bar_right,
230            ))
231
232        # Bottom bar.
233        bar_up = inner_up
234        bar_down = min(image.height - 1, box.down)
235
236        if 0 <= bar_up < image.height and bar_left <= bar_right:
237            hori_bars.append(Box(
238                up=bar_up,
239                down=bar_down,
240                left=bar_left,
241                right=bar_right,
242            ))
243
244    # Render.
245    mask_vert = Mask.from_shapable(image)
246
247    with mask_vert.writable_context:
248        for bar in vert_bars:
249            mask_vert.mat[bar.up:bar.down + 1, bar.left:bar.right + 1] = 1
250
251    fill_vert_dash_gap(
252        dash_thickness=config.dash_thickness,
253        dash_gap=config.dash_gap,
254        mask=mask_vert,
255    )
256
257    mask_hori = Mask.from_shapable(image)
258
259    with mask_hori.writable_context:
260        for bar in hori_bars:
261            mask_hori.mat[bar.up:bar.down + 1, bar.left:bar.right + 1] = 1
262
263    fill_hori_dash_gap(
264        dash_thickness=config.dash_thickness,
265        dash_gap=config.dash_gap,
266        mask=mask_hori,
267    )
268
269    image = image.copy()
270    mask_vert.fill_image(image, config.color, alpha=config.alpha)
271    mask_hori.fill_image(image, config.color, alpha=config.alpha)
272    return image
273
274
275rectangle_streak = Distortion(
276    config_cls=RectangleStreakConfig,
277    state_cls=DistortionNopState[RectangleStreakConfig],
278    func_image=rectangle_streak_image,
279)
280
281
282@attrs.define
283class EllipseStreakConfig(DistortionConfig):
284    thickness: int = 1
285    aspect_ratio: Optional[float] = None
286    short_side_min: int = 10
287    short_side_step: int = 10
288    color: Tuple[int, int, int] = (0, 0, 0)
289    alpha: float = 1.0
290
291
292def ellipse_streak_image(
293    config: EllipseStreakConfig,
294    state: Optional[DistortionNopState[EllipseStreakConfig]],
295    image: Image,
296    rng: Optional[RandomGenerator],
297):
298    mask = Mask.from_shapable(image)
299
300    aspect_ratio = config.aspect_ratio
301    if aspect_ratio is None:
302        aspect_ratio = image.width / image.height
303
304    boxes = generate_centered_boxes(
305        height=image.height,
306        width=image.width,
307        aspect_ratio=aspect_ratio,
308        short_side_min=config.short_side_min,
309        short_side_step=config.short_side_step,
310    )
311    center_y = image.height // 2
312    center_x = image.width // 2
313    center = (center_x, center_y)
314    for box in boxes:
315        mask.assign_mat(
316            cv.ellipse(
317                mask.mat,
318                center=center,
319                axes=(box.width // 2, box.height // 2),
320                angle=0,
321                startAngle=0,
322                endAngle=360,
323                color=1,
324                thickness=config.thickness,
325            )
326        )
327
328    image = image.copy()
329    mask.fill_image(image, config.color, alpha=config.alpha)
330    return image
331
332
333ellipse_streak = Distortion(
334    config_cls=EllipseStreakConfig,
335    state_cls=DistortionNopState[EllipseStreakConfig],
336    func_image=ellipse_streak_image,
337)
def fill_vert_dash_gap(dash_thickness: int, dash_gap: int, mask: vkit.element.mask.Mask):
25def fill_vert_dash_gap(dash_thickness: int, dash_gap: int, mask: Mask):
26    if dash_thickness <= 0 or dash_gap <= 0:
27        return
28
29    with mask.writable_context:
30        step = dash_thickness + dash_gap
31        for offset_y in range(dash_gap):
32            mask.mat[offset_y::step] = 0
def fill_hori_dash_gap(dash_thickness: int, dash_gap: int, mask: vkit.element.mask.Mask):
35def fill_hori_dash_gap(dash_thickness: int, dash_gap: int, mask: Mask):
36    if dash_thickness <= 0 or dash_gap <= 0:
37        return
38
39    with mask.writable_context:
40        step = dash_thickness + dash_gap
41        for offset_x in range(dash_gap):
42            mask.mat[:, offset_x::step] = 0
class LineStreakConfig(vkit.mechanism.distortion.interface.DistortionConfig):
46class LineStreakConfig(DistortionConfig):
47    thickness: int = 1
48    gap: int = 4
49    dash_thickness: int = 0
50    dash_gap: int = 0
51    color: Tuple[int, int, int] = (0, 0, 0)
52    alpha: float = 1.0
53    enable_vert: bool = True
54    enable_hori: bool = True
LineStreakConfig( thickness: int = 1, gap: int = 4, dash_thickness: int = 0, dash_gap: int = 0, color: Tuple[int, int, int] = (0, 0, 0), alpha: float = 1.0, enable_vert: bool = True, enable_hori: bool = True)
 2def __init__(self, thickness=attr_dict['thickness'].default, gap=attr_dict['gap'].default, dash_thickness=attr_dict['dash_thickness'].default, dash_gap=attr_dict['dash_gap'].default, color=attr_dict['color'].default, alpha=attr_dict['alpha'].default, enable_vert=attr_dict['enable_vert'].default, enable_hori=attr_dict['enable_hori'].default):
 3    self.thickness = thickness
 4    self.gap = gap
 5    self.dash_thickness = dash_thickness
 6    self.dash_gap = dash_gap
 7    self.color = color
 8    self.alpha = alpha
 9    self.enable_vert = enable_vert
10    self.enable_hori = enable_hori

Method generated by attrs for class LineStreakConfig.

def line_streak_image( config: vkit.mechanism.distortion.photometric.streak.LineStreakConfig, state: Union[vkit.mechanism.distortion.interface.DistortionNopState[vkit.mechanism.distortion.photometric.streak.LineStreakConfig], NoneType], image: vkit.element.image.Image, rng: Union[numpy.random._generator.Generator, NoneType]):
 57def line_streak_image(
 58    config: LineStreakConfig,
 59    state: Optional[DistortionNopState[LineStreakConfig]],
 60    image: Image,
 61    rng: Optional[RandomGenerator],
 62):
 63    masks: List[Mask] = []
 64
 65    step = config.thickness + config.gap
 66
 67    if config.enable_vert:
 68        mask = Mask.from_shapable(image)
 69
 70        with mask.writable_context:
 71            for offset_x in range(config.thickness):
 72                mask.mat[:, offset_x::step] = 1
 73
 74        fill_vert_dash_gap(
 75            dash_thickness=config.dash_thickness,
 76            dash_gap=config.dash_gap,
 77            mask=mask,
 78        )
 79
 80        masks.append(mask)
 81
 82    if config.enable_hori:
 83        mask = Mask.from_shapable(image)
 84
 85        with mask.writable_context:
 86            for offset_y in range(config.thickness):
 87                mask.mat[offset_y::step] = 1
 88
 89        fill_hori_dash_gap(
 90            dash_thickness=config.dash_thickness,
 91            dash_gap=config.dash_gap,
 92            mask=mask,
 93        )
 94
 95        masks.append(mask)
 96
 97    image = image.copy()
 98    for mask in masks:
 99        mask.fill_image(image, config.color, alpha=config.alpha)
100    return image
def generate_centered_boxes( height: int, width: int, aspect_ratio: float, short_side_min: int, short_side_step: int):
110def generate_centered_boxes(
111    height: int,
112    width: int,
113    aspect_ratio: float,
114    short_side_min: int,
115    short_side_step: int,
116):
117    center_y = height // 2
118    center_x = width // 2
119
120    boxes: List[Box] = []
121    idx = 0
122    while True:
123        short_side = short_side_min + idx * short_side_step
124        if aspect_ratio >= 1:
125            # hori side is longer.
126            height_min = short_side
127            width_min = round(height_min * aspect_ratio)
128        elif 0 < aspect_ratio < 1:
129            # vert side is longer.
130            width_min = short_side
131            height_min = round(width_min / aspect_ratio)
132        else:
133            raise NotImplementedError()
134
135        up = center_y - height_min // 2
136        down = up + height_min - 1
137        left = center_x - width_min // 2
138        right = left + width_min - 1
139
140        if (0 <= up and down < height) or (0 <= left and right < width):
141            boxes.append(Box(up=up, down=down, left=left, right=right))
142            idx += 1
143        else:
144            break
145
146    return boxes
class RectangleStreakConfig(vkit.mechanism.distortion.interface.DistortionConfig):
150class RectangleStreakConfig(DistortionConfig):
151    thickness: int = 1
152    aspect_ratio: Optional[float] = None
153    dash_thickness: int = 0
154    dash_gap: int = 0
155    short_side_min: int = 10
156    short_side_step: int = 10
157    color: Tuple[int, int, int] = (0, 0, 0)
158    alpha: float = 1.0
RectangleStreakConfig( thickness: int = 1, aspect_ratio: Union[float, NoneType] = None, dash_thickness: int = 0, dash_gap: int = 0, short_side_min: int = 10, short_side_step: int = 10, color: Tuple[int, int, int] = (0, 0, 0), alpha: float = 1.0)
 2def __init__(self, thickness=attr_dict['thickness'].default, aspect_ratio=attr_dict['aspect_ratio'].default, dash_thickness=attr_dict['dash_thickness'].default, dash_gap=attr_dict['dash_gap'].default, short_side_min=attr_dict['short_side_min'].default, short_side_step=attr_dict['short_side_step'].default, color=attr_dict['color'].default, alpha=attr_dict['alpha'].default):
 3    self.thickness = thickness
 4    self.aspect_ratio = aspect_ratio
 5    self.dash_thickness = dash_thickness
 6    self.dash_gap = dash_gap
 7    self.short_side_min = short_side_min
 8    self.short_side_step = short_side_step
 9    self.color = color
10    self.alpha = alpha

Method generated by attrs for class RectangleStreakConfig.

def rectangle_streak_image( config: vkit.mechanism.distortion.photometric.streak.RectangleStreakConfig, state: Union[vkit.mechanism.distortion.interface.DistortionNopState[vkit.mechanism.distortion.photometric.streak.RectangleStreakConfig], NoneType], image: vkit.element.image.Image, rng: Union[numpy.random._generator.Generator, NoneType]):
161def rectangle_streak_image(
162    config: RectangleStreakConfig,
163    state: Optional[DistortionNopState[RectangleStreakConfig]],
164    image: Image,
165    rng: Optional[RandomGenerator],
166):
167    aspect_ratio = config.aspect_ratio
168    if aspect_ratio is None:
169        aspect_ratio = image.width / image.height
170
171    boxes = generate_centered_boxes(
172        height=image.height,
173        width=image.width,
174        aspect_ratio=aspect_ratio,
175        short_side_min=config.short_side_min,
176        short_side_step=config.short_side_step,
177    )
178
179    # Generate bars.
180    vert_bars: List[Box] = []
181    hori_bars: List[Box] = []
182
183    for box in boxes:
184        inner_up = box.down - config.thickness + 1
185        inner_down = box.up + config.thickness - 1
186        inner_left = box.right - config.thickness + 1
187        inner_right = box.left + config.thickness - 1
188
189        # Shared by left/right bars.
190        bar_up = max(0, box.up)
191        bar_down = min(image.height - 1, box.down)
192
193        # Left bar.
194        bar_left = max(0, box.left)
195        bar_right = inner_right
196
197        if 0 <= bar_right < image.width and bar_up <= bar_down:
198            vert_bars.append(Box(
199                up=bar_up,
200                down=bar_down,
201                left=bar_left,
202                right=bar_right,
203            ))
204
205        # Right bar.
206        bar_left = inner_left
207        bar_right = min(image.width - 1, box.right)
208
209        if 0 <= bar_left < image.width and bar_up <= bar_down:
210            vert_bars.append(Box(
211                up=bar_up,
212                down=bar_down,
213                left=bar_left,
214                right=bar_right,
215            ))
216
217        # Shared by top/bottom bars.
218        bar_left = max(0, inner_right + 1)
219        bar_right = min(image.width - 1, inner_left - 1)
220
221        # Top bar.
222        bar_up = max(0, box.up)
223        bar_down = inner_down
224
225        if 0 <= inner_down < image.height and bar_left <= bar_right:
226            hori_bars.append(Box(
227                up=bar_up,
228                down=bar_down,
229                left=bar_left,
230                right=bar_right,
231            ))
232
233        # Bottom bar.
234        bar_up = inner_up
235        bar_down = min(image.height - 1, box.down)
236
237        if 0 <= bar_up < image.height and bar_left <= bar_right:
238            hori_bars.append(Box(
239                up=bar_up,
240                down=bar_down,
241                left=bar_left,
242                right=bar_right,
243            ))
244
245    # Render.
246    mask_vert = Mask.from_shapable(image)
247
248    with mask_vert.writable_context:
249        for bar in vert_bars:
250            mask_vert.mat[bar.up:bar.down + 1, bar.left:bar.right + 1] = 1
251
252    fill_vert_dash_gap(
253        dash_thickness=config.dash_thickness,
254        dash_gap=config.dash_gap,
255        mask=mask_vert,
256    )
257
258    mask_hori = Mask.from_shapable(image)
259
260    with mask_hori.writable_context:
261        for bar in hori_bars:
262            mask_hori.mat[bar.up:bar.down + 1, bar.left:bar.right + 1] = 1
263
264    fill_hori_dash_gap(
265        dash_thickness=config.dash_thickness,
266        dash_gap=config.dash_gap,
267        mask=mask_hori,
268    )
269
270    image = image.copy()
271    mask_vert.fill_image(image, config.color, alpha=config.alpha)
272    mask_hori.fill_image(image, config.color, alpha=config.alpha)
273    return image
class EllipseStreakConfig(vkit.mechanism.distortion.interface.DistortionConfig):
284class EllipseStreakConfig(DistortionConfig):
285    thickness: int = 1
286    aspect_ratio: Optional[float] = None
287    short_side_min: int = 10
288    short_side_step: int = 10
289    color: Tuple[int, int, int] = (0, 0, 0)
290    alpha: float = 1.0
EllipseStreakConfig( thickness: int = 1, aspect_ratio: Union[float, NoneType] = None, short_side_min: int = 10, short_side_step: int = 10, color: Tuple[int, int, int] = (0, 0, 0), alpha: float = 1.0)
2def __init__(self, thickness=attr_dict['thickness'].default, aspect_ratio=attr_dict['aspect_ratio'].default, short_side_min=attr_dict['short_side_min'].default, short_side_step=attr_dict['short_side_step'].default, color=attr_dict['color'].default, alpha=attr_dict['alpha'].default):
3    self.thickness = thickness
4    self.aspect_ratio = aspect_ratio
5    self.short_side_min = short_side_min
6    self.short_side_step = short_side_step
7    self.color = color
8    self.alpha = alpha

Method generated by attrs for class EllipseStreakConfig.

def ellipse_streak_image( config: vkit.mechanism.distortion.photometric.streak.EllipseStreakConfig, state: Union[vkit.mechanism.distortion.interface.DistortionNopState[vkit.mechanism.distortion.photometric.streak.EllipseStreakConfig], NoneType], image: vkit.element.image.Image, rng: Union[numpy.random._generator.Generator, NoneType]):
293def ellipse_streak_image(
294    config: EllipseStreakConfig,
295    state: Optional[DistortionNopState[EllipseStreakConfig]],
296    image: Image,
297    rng: Optional[RandomGenerator],
298):
299    mask = Mask.from_shapable(image)
300
301    aspect_ratio = config.aspect_ratio
302    if aspect_ratio is None:
303        aspect_ratio = image.width / image.height
304
305    boxes = generate_centered_boxes(
306        height=image.height,
307        width=image.width,
308        aspect_ratio=aspect_ratio,
309        short_side_min=config.short_side_min,
310        short_side_step=config.short_side_step,
311    )
312    center_y = image.height // 2
313    center_x = image.width // 2
314    center = (center_x, center_y)
315    for box in boxes:
316        mask.assign_mat(
317            cv.ellipse(
318                mask.mat,
319                center=center,
320                axes=(box.width // 2, box.height // 2),
321                angle=0,
322                startAngle=0,
323                endAngle=360,
324                color=1,
325                thickness=config.thickness,
326            )
327        )
328
329    image = image.copy()
330    mask.fill_image(image, config.color, alpha=config.alpha)
331    return image