diff options
Diffstat (limited to 'text_recognizer/data/iam_paragraphs.py')
| -rw-r--r-- | text_recognizer/data/iam_paragraphs.py | 4 |
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) |