summaryrefslogtreecommitdiff
path: root/text_recognizer/data
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2022-09-27 23:18:06 +0200
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2022-09-27 23:18:06 +0200
commit2dcbd13eb311d0a4b31e086a53d2afc55d50e0b5 (patch)
tree29a290313f0967177391265a3b78c1ba055ccffa /text_recognizer/data
parent626cf824526b00b18043315322f86789278939fa (diff)
Add typing
Diffstat (limited to 'text_recognizer/data')
-rw-r--r--text_recognizer/data/stems/line.py15
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):