Skip to content

Commit 1992c0c

Browse files
BelvalBelval
committed
Add tests
1 parent 942e4f4 commit 1992c0c

6 files changed

Lines changed: 141 additions & 14 deletions

File tree

‎TextRecognitionDataGenerator/data_generator.py‎

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,18 @@ def generate_from_tuple(cls, t):
2121
cls.generate(*t)
2222

2323
@classmethod
24-
def generate(cls, index, text, font, out_dir, height, extension, skewing_angle, random_skew, blur, random_blur, background_type, distorsion_type, distorsion_orientation, is_handwritten, name_format, width, alignment, text_color, orientation, space_width):
24+
def generate(cls, index, text, font, out_dir, size, extension, skewing_angle, random_skew, blur, random_blur, background_type, distorsion_type, distorsion_orientation, is_handwritten, name_format, width, alignment, text_color, orientation, space_width):
2525
image = None
2626

2727
##########################
2828
# Create picture of text #
2929
##########################
3030
if is_handwritten:
3131
if orientation == 1:
32-
raise Exception("Vertical handwritten text is unavailable")
32+
raise ValueError("Vertical handwritten text is unavailable")
3333
image = HandwrittenTextGenerator.generate(text)
3434
else:
35-
image = ComputerTextGenerator.generate(text, font, text_color, height, orientation, space_width)
35+
image = ComputerTextGenerator.generate(text, font, text_color, size, orientation, space_width)
3636

3737
random_angle = random.randint(0-skewing_angle, skewing_angle)
3838

@@ -66,23 +66,32 @@ def generate(cls, index, text, font, out_dir, height, extension, skewing_angle,
6666
# Resize image to desired format #
6767
##################################
6868

69-
new_width = int(float(distorted_img.size[0] + 10) * (float(height) / float(distorted_img.size[1] + 10)))
70-
71-
resized_img = distorted_img.resize((new_width, height - 10), Image.ANTIALIAS)
72-
73-
background_width = width if width > 0 else new_width + 10
69+
# Horizontal text
70+
if orientation == 0:
71+
new_width = int(float(distorted_img.size[0] + 10) * (float(size) / float(distorted_img.size[1] + 10)))
72+
resized_img = distorted_img.resize((new_width, size - 10), Image.ANTIALIAS)
73+
background_width = width if width > 0 else new_width + 10
74+
background_height = size
75+
# Vertical text
76+
elif orientation == 1:
77+
new_height = int(float(distorted_img.size[1] + 10) * (float(size) / float(distorted_img.size[0] + 10)))
78+
resized_img = distorted_img.resize((size - 10, new_height), Image.ANTIALIAS)
79+
background_width = size
80+
background_height = new_height + 10
81+
else:
82+
raise ValueError("Invalid orientation")
7483

7584
#############################
7685
# Generate background image #
7786
#############################
7887
if background_type == 0:
79-
background = BackgroundGenerator.gaussian_noise(height, background_width)
88+
background = BackgroundGenerator.gaussian_noise(background_height, background_width)
8089
elif background_type == 1:
81-
background = BackgroundGenerator.plain_white(height, background_width)
90+
background = BackgroundGenerator.plain_white(background_height, background_width)
8291
elif background_type == 2:
83-
background = BackgroundGenerator.quasicrystal(height, background_width)
92+
background = BackgroundGenerator.quasicrystal(background_height, background_width)
8493
else:
85-
background = BackgroundGenerator.picture(height, background_width)
94+
background = BackgroundGenerator.picture(background_height, background_width)
8695

8796
#############################
8897
# Place text with alignment #

‎TextRecognitionDataGenerator/run.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ def parse_arguments():
103103
"--format",
104104
type=int,
105105
nargs="?",
106-
help="Define the height of the produced images",
106+
help="Define the height of the produced images if horizontal, else the width",
107107
default=32,
108108
)
109109
parser.add_argument(
@@ -218,7 +218,7 @@ def parse_arguments():
218218
type=int,
219219
nargs="?",
220220
help="Define the orientation of the text. 0: Horizontal, 1: Vertical",
221-
default=1
221+
default=0
222222
)
223223
parser.add_argument(
224224
"-tc",

‎tests.py‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,124 @@ def test_generate_data_with_right_alignment(self):
335335

336336
os.remove('tests/out/TEST TEST TEST_8.jpg')
337337

338+
def test_raise_if_handwritten_and_vertical(self):
339+
try:
340+
FakeTextDataGenerator.generate(
341+
9,
342+
'TEST TEST TEST',
343+
'tests/font.ttf',
344+
'tests/out/',
345+
64,
346+
'jpg',
347+
0,
348+
False,
349+
0,
350+
False,
351+
1,
352+
0,
353+
0,
354+
True,
355+
0,
356+
1000,
357+
2,
358+
'#010101',
359+
1,
360+
1
361+
)
362+
raise Exception("Vertical handwritten did not throw")
363+
except ValueError:
364+
pass
365+
366+
def test_generate_vertical_text(self):
367+
FakeTextDataGenerator.generate(
368+
10,
369+
'TEST TEST TEST',
370+
'tests/font.ttf',
371+
'tests/out/',
372+
32,
373+
'jpg',
374+
0,
375+
False,
376+
0,
377+
False,
378+
1,
379+
0,
380+
0,
381+
False,
382+
0,
383+
-1,
384+
0,
385+
'#010101',
386+
1,
387+
1
388+
)
389+
390+
self.assertTrue(
391+
md5('tests/out/TEST TEST TEST_10.jpg') == md5('tests/expected_results/TEST TEST TEST_10.jpg')
392+
)
393+
394+
os.remove('tests/out/TEST TEST TEST_10.jpg')
395+
396+
def test_generate_horizontal_text_with_variable_space(self):
397+
FakeTextDataGenerator.generate(
398+
11,
399+
'TEST TEST TEST',
400+
'tests/font.ttf',
401+
'tests/out/',
402+
32,
403+
'jpg',
404+
0,
405+
False,
406+
0,
407+
False,
408+
1,
409+
0,
410+
0,
411+
False,
412+
0,
413+
-1,
414+
0,
415+
'#010101',
416+
0,
417+
4
418+
)
419+
420+
self.assertTrue(
421+
md5('tests/out/TEST TEST TEST_11.jpg') == md5('tests/expected_results/TEST TEST TEST_11.jpg')
422+
)
423+
424+
os.remove('tests/out/TEST TEST TEST_11.jpg')
425+
426+
def test_generate_vertical_text_with_variable_space(self):
427+
FakeTextDataGenerator.generate(
428+
12,
429+
'TEST TEST TEST',
430+
'tests/font.ttf',
431+
'tests/out/',
432+
32,
433+
'jpg',
434+
0,
435+
False,
436+
0,
437+
False,
438+
1,
439+
0,
440+
0,
441+
False,
442+
0,
443+
-1,
444+
0,
445+
'#010101',
446+
1,
447+
2
448+
)
449+
450+
self.assertTrue(
451+
md5('tests/out/TEST TEST TEST_12.jpg') == md5('tests/expected_results/TEST TEST TEST_12.jpg')
452+
)
453+
454+
os.remove('tests/out/TEST TEST TEST_12.jpg')
455+
338456
def test_generate_string_with_letters(self):
339457
s = create_strings_randomly(1, False, 1, True, False, False, 'en')[0]
340458

‎tests/expected_results/TEST TEST TEST_10.jpg‎

3.71 KB
Loading

‎tests/expected_results/TEST TEST TEST_11.jpg‎

2.35 KB
Loading

‎tests/expected_results/TEST TEST TEST_12.jpg‎

3.74 KB
Loading

0 commit comments

Comments
 (0)