summaryrefslogtreecommitdiff
path: root/text_recognizer/data/iam_paragraphs.py
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2021-06-27 20:25:25 +0200
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2021-06-27 20:25:25 +0200
commit442eac315e4b8be19adab80fb7332d29f68c077c (patch)
tree1084ff180e7852029918534343ec2a18b6b0485f /text_recognizer/data/iam_paragraphs.py
parentcafd6b8b10d804b3eee235652cb5218ef4a469b4 (diff)
Fixed bug in word pieces
Diffstat (limited to 'text_recognizer/data/iam_paragraphs.py')
-rw-r--r--text_recognizer/data/iam_paragraphs.py4
1 files changed, 4 insertions, 0 deletions
diff --git a/text_recognizer/data/iam_paragraphs.py b/text_recognizer/data/iam_paragraphs.py
index 6022804..fe60e99 100644
--- a/text_recognizer/data/iam_paragraphs.py
+++ b/text_recognizer/data/iam_paragraphs.py
@@ -17,6 +17,7 @@ from text_recognizer.data.base_dataset import (
)
from text_recognizer.data.base_data_module import BaseDataModule, load_and_print_info
from text_recognizer.data.emnist import emnist_mapping
+from text_recognizer.data.mappings import WordPieceMapping
from text_recognizer.data.iam import IAM
from text_recognizer.data.transforms import WordPiece
@@ -49,6 +50,9 @@ class IAMParagraphs(BaseDataModule):
self.mapping, self.inverse_mapping, _ = emnist_mapping(
extra_symbols=[NEW_LINE_TOKEN]
)
+ if word_pieces:
+ self.mapping = WordPieceMapping()
+
self.train_fraction = train_fraction
self.dims = (1, IMAGE_HEIGHT, IMAGE_WIDTH)