diff options
Diffstat (limited to 'text_recognizer/data/iam_synthetic_paragraphs.py')
| -rw-r--r-- | text_recognizer/data/iam_synthetic_paragraphs.py | 7 |
1 files changed, 5 insertions, 2 deletions
diff --git a/text_recognizer/data/iam_synthetic_paragraphs.py b/text_recognizer/data/iam_synthetic_paragraphs.py index 9f1bd12..4ccc5c2 100644 --- a/text_recognizer/data/iam_synthetic_paragraphs.py +++ b/text_recognizer/data/iam_synthetic_paragraphs.py @@ -18,6 +18,7 @@ from text_recognizer.data.base_data_module import BaseDataModule, load_and_print from text_recognizer.data.iam_paragraphs import ( get_dataset_properties, get_transform, + get_target_transform, NEW_LINE_TOKEN, IAMParagraphs, IMAGE_SCALE_FACTOR, @@ -41,12 +42,13 @@ class IAMSyntheticParagraphs(IAMParagraphs): def __init__( self, - batch_size: int = 128, + batch_size: int = 16, num_workers: int = 0, train_fraction: float = 0.8, augment: bool = True, + word_pieces: bool = False, ) -> None: - super().__init__(batch_size, num_workers, train_fraction, augment) + super().__init__(batch_size, num_workers, train_fraction, augment, word_pieces) def prepare_data(self) -> None: """Prepare IAM lines to be used to generate paragraphs.""" @@ -95,6 +97,7 @@ class IAMSyntheticParagraphs(IAMParagraphs): transform=get_transform( image_shape=self.dims[1:], augment=self.augment ), + target_transform=get_target_transform(self.word_pieces) ) def __repr__(self) -> str: |