From 25b5d6983d51e0e791b96a76beb7e49f392cd9a8 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Mon, 7 Dec 2020 22:54:04 +0100 Subject: Segmentation working! --- src/text_recognizer/datasets/emnist_lines_dataset.py | 11 +++++++++++ 1 file changed, 11 insertions(+) (limited to 'src/text_recognizer/datasets/emnist_lines_dataset.py') diff --git a/src/text_recognizer/datasets/emnist_lines_dataset.py b/src/text_recognizer/datasets/emnist_lines_dataset.py index 6871492..eddf341 100644 --- a/src/text_recognizer/datasets/emnist_lines_dataset.py +++ b/src/text_recognizer/datasets/emnist_lines_dataset.py @@ -10,6 +10,7 @@ from loguru import logger import numpy as np import torch from torch import Tensor +import torch.nn.functional as F from torchvision.transforms import ToTensor from text_recognizer.datasets.dataset import Dataset @@ -23,6 +24,8 @@ from text_recognizer.datasets.util import ( DATA_DIRNAME = DATA_DIRNAME / "processed" / "emnist_lines" +MAX_WIDTH = 952 + class EmnistLinesDataset(Dataset): """Synthetic dataset of lines from the Brown corpus with Emnist characters.""" @@ -254,6 +257,14 @@ def construct_image_from_string( for image in sampled_images: concatenated_image[:, x : (x + width)] += image x += next_overlap_width + + if concatenated_image.shape[-1] > MAX_WIDTH: + concatenated_image = Tensor(concatenated_image).unsqueeze(0) + concatenated_image = F.interpolate( + concatenated_image, size=MAX_WIDTH, mode="nearest" + ) + concatenated_image = concatenated_image.squeeze(0).numpy() + return np.minimum(255, concatenated_image) -- cgit v1.2.3-70-g09d2