forked from Belval/TextRecognitionDataGenerator
-
Notifications
You must be signed in to change notification settings - Fork 0
/
from_random.py
121 lines (113 loc) · 3.59 KB
/
from_random.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
114
115
116
117
118
119
120
121
import os
from typing import List, Tuple
from trdg.generators.from_strings import GeneratorFromStrings
from trdg.data_generator import FakeTextDataGenerator
from trdg.string_generator import create_strings_randomly
from trdg.utils import load_dict, load_fonts
class GeneratorFromRandom:
"""Generator that uses randomly generated words"""
def __init__(
self,
count: int = -1,
length: int = 1,
allow_variable: bool = False,
use_letters: bool = True,
use_numbers: bool = True,
use_symbols: bool = True,
fonts: List[str] = [],
language: str = "en",
size: int = 32,
skewing_angle: int = 0,
random_skew: bool = False,
blur: int = 0,
random_blur: bool = False,
background_type: int = 0,
distorsion_type: int = 0,
distorsion_orientation: int = 0,
is_handwritten: bool = False,
width: int = -1,
alignment: int = 1,
text_color: str = "#282828",
orientation: int = 0,
space_width: float = 1.0,
character_spacing: int = 0,
margins: Tuple[int, int, int, int] = (5, 5, 5, 5),
fit: bool = False,
output_mask: bool = False,
word_split: bool = False,
image_dir: str = os.path.join(
"..", os.path.split(os.path.realpath(__file__))[0], "images"
),
stroke_width: int = 0,
stroke_fill: str = "#282828",
image_mode: str = "RGB",
output_bboxes: int = 0,
):
self.generated_count = 0
self.count = count
self.length = length
self.allow_variable = allow_variable
self.use_letters = use_letters
self.use_numbers = use_numbers
self.use_symbols = use_symbols
self.language = language
self.batch_size = min(max(count, 1), 1000)
self.steps_until_regeneration = self.batch_size
self.generator = GeneratorFromStrings(
create_strings_randomly(
self.length,
self.allow_variable,
self.batch_size,
self.use_letters,
self.use_numbers,
self.use_symbols,
self.language,
),
count,
fonts if len(fonts) else load_fonts(language),
language,
size,
skewing_angle,
random_skew,
blur,
random_blur,
background_type,
distorsion_type,
distorsion_orientation,
is_handwritten,
width,
alignment,
text_color,
orientation,
space_width,
character_spacing,
margins,
fit,
output_mask,
word_split,
image_dir,
stroke_width,
stroke_fill,
image_mode,
output_bboxes,
)
def __iter__(self):
return self
def __next__(self):
if self.generated_count == self.count:
raise StopIteration
self.generated_count += 1
return self.next()
def next(self):
if self.generator.generated_count >= self.steps_until_regeneration:
self.generator.strings = create_strings_randomly(
self.length,
self.allow_variable,
self.batch_size,
self.use_letters,
self.use_numbers,
self.use_symbols,
self.language,
)
self.steps_until_regeneration += self.batch_size
return self.generator.next()