-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
200 lines (157 loc) · 7.52 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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
import argparse
import json
from typing import Tuple, List
import cv2
import editdistance
from path import Path
from dataloader_iam import DataLoaderIAM, Batch
from model import Model, DecoderType
from preprocessor import Preprocessor
class FilePaths:
"""Filenames and paths to data."""
fn_char_list = '../model/charList.txt'
fn_summary = '../model/summary.json'
fn_corpus = '../data/corpus.txt'
def get_img_height() -> int:
"""Fixed height for NN."""
return 32
def get_img_size(line_mode: bool = False) -> Tuple[int, int]:
"""Height is fixed for NN, width is set according to training mode (single words or text lines)."""
if line_mode:
return 256, get_img_height()
return 128, get_img_height()
def write_summary(char_error_rates: List[float], word_accuracies: List[float]) -> None:
"""Writes training summary file for NN."""
with open(FilePaths.fn_summary, 'w') as f:
json.dump({'charErrorRates': char_error_rates, 'wordAccuracies': word_accuracies}, f)
def char_list_from_file() -> List[str]:
with open(FilePaths.fn_char_list) as f:
return list(f.read())
def train(model: Model,
loader: DataLoaderIAM,
line_mode: bool,
early_stopping: int = 25) -> None:
"""Trains NN."""
epoch = 0 # number of training epochs since start
summary_char_error_rates = []
summary_word_accuracies = []
preprocessor = Preprocessor(get_img_size(line_mode), data_augmentation=True, line_mode=line_mode)
best_char_error_rate = float('inf') # best validation character error rate
no_improvement_since = 0 # number of epochs no improvement of character error rate occurred
# stop training after this number of epochs without improvement
while True:
epoch += 1
print('Epoch:', epoch)
# train
print('Train NN')
loader.train_set()
while loader.has_next():
iter_info = loader.get_iterator_info()
batch = loader.get_next()
batch = preprocessor.process_batch(batch)
loss = model.train_batch(batch)
print(f'Epoch: {epoch} Batch: {iter_info[0]}/{iter_info[1]} Loss: {loss}')
# validate
char_error_rate, word_accuracy = validate(model, loader, line_mode)
# write summary
summary_char_error_rates.append(char_error_rate)
summary_word_accuracies.append(word_accuracy)
write_summary(summary_char_error_rates, summary_word_accuracies)
# if best validation accuracy so far, save model parameters
if char_error_rate < best_char_error_rate:
print('Character error rate improved, save model')
best_char_error_rate = char_error_rate
no_improvement_since = 0
model.save()
else:
print(f'Character error rate not improved, best so far: {char_error_rate * 100.0}%')
no_improvement_since += 1
# stop training if no more improvement in the last x epochs
if no_improvement_since >= early_stopping:
print(f'No more improvement since {early_stopping} epochs. Training stopped.')
break
def validate(model: Model, loader: DataLoaderIAM, line_mode: bool) -> Tuple[float, float]:
"""Validates NN."""
print('Validate NN')
loader.validation_set()
preprocessor = Preprocessor(get_img_size(line_mode), line_mode=line_mode)
num_char_err = 0
num_char_total = 0
num_word_ok = 0
num_word_total = 0
while loader.has_next():
iter_info = loader.get_iterator_info()
print(f'Batch: {iter_info[0]} / {iter_info[1]}')
batch = loader.get_next()
batch = preprocessor.process_batch(batch)
recognized, _ = model.infer_batch(batch)
print('Ground truth -> Recognized')
for i in range(len(recognized)):
num_word_ok += 1 if batch.gt_texts[i] == recognized[i] else 0
num_word_total += 1
dist = editdistance.eval(recognized[i], batch.gt_texts[i])
num_char_err += dist
num_char_total += len(batch.gt_texts[i])
print('[OK]' if dist == 0 else '[ERR:%d]' % dist, '"' + batch.gt_texts[i] + '"', '->',
'"' + recognized[i] + '"')
# print validation result
char_error_rate = num_char_err / num_char_total
word_accuracy = num_word_ok / num_word_total
print(f'Character error rate: {char_error_rate * 100.0}%. Word accuracy: {word_accuracy * 100.0}%.')
return char_error_rate, word_accuracy
def infer(model: Model, fn_img: Path) -> None:
"""Recognizes text in image provided by file path."""
img = cv2.imread(fn_img, cv2.IMREAD_GRAYSCALE)
assert img is not None
preprocessor = Preprocessor(get_img_size(), dynamic_width=True, padding=16)
img = preprocessor.process_img(img)
batch = Batch([img], None, 1)
recognized, probability = model.infer_batch(batch, True)
print(f'Recognized: "{recognized[0]}"')
print(f'Probability: {probability[0]}')
def parse_args() -> argparse.Namespace:
"""Parses arguments from the command line."""
parser = argparse.ArgumentParser()
parser.add_argument('--mode', choices=['train', 'validate', 'infer'], default='infer')
parser.add_argument('--decoder', choices=['bestpath', 'beamsearch', 'wordbeamsearch'], default='bestpath')
parser.add_argument('--batch_size', help='Batch size.', type=int, default=100)
parser.add_argument('--data_dir', help='Directory containing IAM dataset.', type=Path, required=False)
parser.add_argument('--fast', help='Load samples from LMDB.', action='store_true')
parser.add_argument('--line_mode', help='Train to read text lines instead of single words.', action='store_true')
parser.add_argument('--img_file', help='Image used for inference.', type=Path, default='../data/word.png')
parser.add_argument('--early_stopping', help='Early stopping epochs.', type=int, default=25)
parser.add_argument('--dump', help='Dump output of NN to CSV file(s).', action='store_true')
return parser.parse_args()
def main():
"""Main function."""
# parse arguments and set CTC decoder
args = parse_args()
decoder_mapping = {'bestpath': DecoderType.BestPath,
'beamsearch': DecoderType.BeamSearch,
'wordbeamsearch': DecoderType.WordBeamSearch}
decoder_type = decoder_mapping[args.decoder]
# train the model
if args.mode == 'train':
loader = DataLoaderIAM(args.data_dir, args.batch_size, fast=args.fast)
# when in line mode, take care to have a whitespace in the char list
char_list = loader.char_list
if args.line_mode and ' ' not in char_list:
char_list = [' '] + char_list
# save characters and words
with open(FilePaths.fn_char_list, 'w') as f:
f.write(''.join(char_list))
with open(FilePaths.fn_corpus, 'w') as f:
f.write(' '.join(loader.train_words + loader.validation_words))
model = Model(char_list, decoder_type)
train(model, loader, line_mode=args.line_mode, early_stopping=args.early_stopping)
# evaluate it on the validation set
elif args.mode == 'validate':
loader = DataLoaderIAM(args.data_dir, args.batch_size, fast=args.fast)
model = Model(char_list_from_file(), decoder_type, must_restore=True)
validate(model, loader, args.line_mode)
# infer text on test image
elif args.mode == 'infer':
model = Model(char_list_from_file(), decoder_type, must_restore=True, dump=args.dump)
infer(model, args.img_file)
if __name__ == '__main__':
main()