diff options
Diffstat (limited to 'text_recognizer/data')
| -rw-r--r-- | text_recognizer/data/stems/line.py | 15 |
1 files changed, 11 insertions, 4 deletions
diff --git a/text_recognizer/data/stems/line.py b/text_recognizer/data/stems/line.py index 2fe1a2c..4f0ce05 100644 --- a/text_recognizer/data/stems/line.py +++ b/text_recognizer/data/stems/line.py @@ -1,4 +1,5 @@ import random +from typing import Any, Dict from PIL import Image import torchvision.transforms as T @@ -11,8 +12,11 @@ class LineStem(ImageStem): """A stem for handling images containing a line of text.""" def __init__( - self, augment=False, color_jitter_kwargs=None, random_affine_kwargs=None - ): + self, + augment: bool = False, + color_jitter_kwargs: Dict[str, Any] = None, + random_affine_kwargs: Dict[str, Any] = None, + ) -> None: super().__init__() if color_jitter_kwargs is None: color_jitter_kwargs = {"brightness": (0.5, 1)} @@ -39,8 +43,11 @@ class IamLinesStem(ImageStem): """A stem for handling images containing lines of text from the IAMLines dataset.""" def __init__( - self, augment=False, color_jitter_kwargs=None, random_affine_kwargs=None - ): + self, + augment: bool = False, + color_jitter_kwargs: Dict[str, Any] = None, + random_affine_kwargs: Dict[str, Any] = None, + ) -> None: super().__init__() def embed_crop(crop, augment=augment): |