/
lowinkperiodiclines.py
113 lines (97 loc) · 4.39 KB
/
lowinkperiodiclines.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import random
from augraphy.augmentations.lowinkline import LowInkLine
class LowInkPeriodicLines(LowInkLine):
"""Creates a set of lines that repeat in a periodic fashion throughout the
image.
:param count_range: Pair of ints determining the range from which to sample
the number of lines to apply.
:type count_range: tuple, optional
:param period_range: Pair of ints determining the range from which to sample
the distance between lines.
:type period_range: tuple, optional
:param use_consistent_lines: Whether or not to vary the width and alpha of
generated low ink lines.
:type use_consistent_lines: bool, optional
:param noise_probability: The probability to add noise into the generated lines.
:type noise_probability: float, optional
:param p: The probability that this Augmentation will be applied.
:type p: float, optional
"""
def __init__(
self,
count_range=(2, 5),
period_range=(10, 30),
use_consistent_lines=True,
noise_probability=0.1,
p=1,
):
"""Constructor method"""
super().__init__(
use_consistent_lines=use_consistent_lines,
noise_probability=noise_probability,
p=p,
)
self.count_range = count_range
self.period_range = period_range
# Constructs a string representation of this Augmentation.
def __repr__(self):
return f"LowInkPeriodicLines(count_range={self.count_range}, period_range={self.period_range}, use_consistent_lines={self.use_consistent_lines}, p={self.p})"
def add_periodic_transparency_line(self, mask, line_count, offset, alpha):
"""Creates horizontal lines of some opacity over the input image, at y-positions determined by the offset and line_count.
:param mask: The image to apply the line to.
:type mask: numpy.array
:param line_count: The number of lines to generate.
:type line_count: int
:param offset: How far from the edge of the image to generate lines.
:type offset: int
:param alpha: The opacity of the lines.
:type alpha: int
"""
period = mask.shape[0] // line_count
for y in range(mask.shape[0] - offset):
if (period != 0) and (
y % period == 0
): # period can't be zero here, else there would be zero division error
self.add_transparency_line(mask, y + offset, alpha)
def add_periodic_transparency_lines(self, mask, lines, line_periods):
"""Creates horizontal lines of random opacity over the input image, at
random intervals.
:param mask: The image to apply the line to.
:type mask: numpy.array
:param lines: How many lines to add to the image.
:type lines: int
:param line_periods: The distance between lines.
:type line_periods: int
"""
period = mask.shape[0] // line_periods
self.add_periodic_transparency_line(
mask,
line_periods,
offset=random.randint(0, 5),
alpha=random.randint(96, 255),
)
for i in range(lines):
self.add_periodic_transparency_line(
mask,
line_periods,
offset=random.randint(0, period),
alpha=random.randint(16, 96),
)
# Applies the Augmentation to input data.
def __call__(self, image, layer=None, mask=None, keypoints=None, bounding_boxes=None, force=False):
if force or self.should_run():
image = image.copy()
count = random.randint(self.count_range[0], self.count_range[1])
period = random.randint(self.period_range[0], self.period_range[1])
for i in range(count):
self.add_periodic_transparency_lines(image, count, period)
# check for additional output of mask, keypoints and bounding boxes
outputs_extra = []
if mask is not None or keypoints is not None or bounding_boxes is not None:
outputs_extra = [mask, keypoints, bounding_boxes]
# returns additional mask, keypoints and bounding boxes if there is additional input
if outputs_extra:
# returns in the format of [image, mask, keypoints, bounding_boxes]
return [image] + outputs_extra
else:
return image