From 3a21c29e2eff4378c63717f8920ca3ccbfef013c Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sun, 3 Oct 2021 00:31:00 +0200 Subject: Lint files --- text_recognizer/data/emnist_lines.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) (limited to 'text_recognizer/data/emnist_lines.py') diff --git a/text_recognizer/data/emnist_lines.py b/text_recognizer/data/emnist_lines.py index d4b2b40..3ff8a54 100644 --- a/text_recognizer/data/emnist_lines.py +++ b/text_recognizer/data/emnist_lines.py @@ -8,6 +8,7 @@ import h5py from loguru import logger as log import numpy as np import torch +from torch import Tensor from torchvision import transforms from torchvision.transforms.functional import InterpolationMode @@ -190,7 +191,9 @@ def _get_samples_by_char( return samples_by_char -def _select_letter_samples_for_string(string: str, samples_by_char: defaultdict): +def _select_letter_samples_for_string( + string: str, samples_by_char: defaultdict +) -> List[Tensor]: null_image = torch.zeros((28, 28), dtype=torch.uint8) sample_image_by_char = {} for char in string: @@ -208,7 +211,7 @@ def _construct_image_from_string( min_overlap: float, max_overlap: float, width: int, -) -> torch.Tensor: +) -> Tensor: overlap = np.random.uniform(min_overlap, max_overlap) sampled_images = _select_letter_samples_for_string(string, samples_by_char) H, W = sampled_images[0].shape @@ -218,7 +221,7 @@ def _construct_image_from_string( for image in sampled_images: concatenated_image[:, x : (x + W)] += image x += next_overlap_width - return torch.minimum(torch.Tensor([255]), concatenated_image) + return torch.minimum(Tensor([255]), concatenated_image) def _create_dataset_of_images( @@ -228,7 +231,7 @@ def _create_dataset_of_images( min_overlap: float, max_overlap: float, dims: Tuple, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[Tensor, Tensor]: images = torch.zeros((num_samples, IMAGE_HEIGHT, dims[2])) labels = [] for n in range(num_samples): -- cgit v1.2.3-70-g09d2