-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain.py
115 lines (91 loc) · 3.54 KB
/
main.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
import time
import numpy as np
from PIL import Image, ImageFilter
from matplotlib import pyplot as plt
from skimage.exposure import histogram
import cnn
import dataset
import processing as prc
def create_dataset(image_name):
im = Image.open(image_name).convert('L')
rotated = prc.get_pretty_rotated(im)
lines_positions = prc.get_lines_positions(rotated)
lines = prc.get_lines_from_positions(rotated, lines_positions)
result_letters = []
for line in lines:
words_positions = prc.get_words_positions(line)
words = prc.get_words_from_position(line, words_positions)
for word in words:
letters_bounds = prc.get_letters_bounds(word)
letters = prc.get_letters_from_bounds(word, letters_bounds)
result_letters += letters
plt.figure(1)
sps = int(len(result_letters) ** .5) + 1
for i in range(len(result_letters)):
plt.subplot(sps, sps, i + 1)
plt.imshow(result_letters[i], cmap='gray')
plt.show()
ds = dataset.Dataset()
real_letters = input('>> ')
for i in range(len(real_letters)):
ds.add_letter(result_letters[i], real_letters[i])
def recognize_text(image_name):
model = cnn.CNN()
im = Image.open(image_name).convert('L')
im = prc.get_prepared(im)
rotated = prc.get_pretty_rotated(im)
lines_positions = prc.get_lines_positions(rotated)
lines = prc.get_lines_from_positions(rotated, lines_positions)
result_letters = []
for line in lines:
pretty_line = prc.get_pretty_sloped(line)
# plt.figure(1)
# plt.subplot(211)
# plt.imshow(line, cmap='gray')
# plt.subplot(212)
# plt.imshow(pretty_line, cmap='gray')
words_positions = prc.get_words_positions(pretty_line)
# for wp in words_positions:
# plt.plot([wp[0], wp[0]], [0, line.size[1]], color='r')
# plt.plot([wp[1], wp[1]], [0, line.size[1]], color='r')
# plt.show()
words = prc.get_words_from_position(pretty_line, words_positions)
for word in words:
letters_bounds = prc.get_letters_bounds(word)
letters = prc.get_letters_from_bounds(word, letters_bounds)
# plt.figure(1)
# i = 1
# l = int(len(letters) ** .5) + 1
# for letter in letters:
# plt.subplot(l, l, i)
# i += 1
# plt.imshow(letter)
# plt.show()
for letter in letters:
guessed_letter, prob = model.predict(letter, return_probs=True)
result_letters.append(guessed_letter)
# plt.figure(1)
# plt.subplot(121)
# plt.imshow(letter, cmap='gray')
# plt.subplot(122)
# for letter, prob in prob.items():
# if prob > .005:
# plt.bar(letter, prob)
# plt.show()
result_letters.append(' ')
result_letters.append('\n')
text = ''.join(result_letters)
text = text.replace('ь|', 'ы')
text = text.replace('ъ|', 'ы')
return text
def retrain():
model = cnn.CNN()
ds = dataset.Dataset()
images, letters = ds.get_data()
model.train_and_save(images, letters)
if __name__ == "__main__":
for file in ['sample_2.png', 'sample_cursive.png', 'sample_rotated_little.png']:
print('------------------', file, '---------------------')
start_time = time.time()
print(recognize_text(file))
print('\n\ntook = %.3f sec' % (time.time() - start_time))