From 2dcbd13eb311d0a4b31e086a53d2afc55d50e0b5 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Tue, 27 Sep 2022 23:18:06 +0200 Subject: Add typing --- text_recognizer/data/stems/line.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) (limited to 'text_recognizer/data/stems') 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): -- cgit v1.2.3-70-g09d2