diff options
Diffstat (limited to 'text_recognizer')
| -rw-r--r-- | text_recognizer/datasets/base_data_module.py | 2 | ||||
| -rw-r--r-- | text_recognizer/datasets/base_dataset.py | 8 | ||||
| -rw-r--r-- | text_recognizer/datasets/emnist.py | 12 | ||||
| -rw-r--r-- | text_recognizer/datasets/emnist_essentials.json | 2 | ||||
| -rw-r--r-- | text_recognizer/datasets/emnist_lines.py | 172 | ||||
| -rw-r--r-- | text_recognizer/datasets/sentence_generator.py | 30 |
6 files changed, 165 insertions, 61 deletions
diff --git a/text_recognizer/datasets/base_data_module.py b/text_recognizer/datasets/base_data_module.py index 830b39b..f5e7300 100644 --- a/text_recognizer/datasets/base_data_module.py +++ b/text_recognizer/datasets/base_data_module.py @@ -46,7 +46,7 @@ class BaseDataModule(pl.LightningDataModule): def setup(self, stage: str = None) -> None: """Split into train, val, test, and set dims. - + Should assign `torch Dataset` objects to self.data_train, self.data_val, and optionally self.data_test. diff --git a/text_recognizer/datasets/base_dataset.py b/text_recognizer/datasets/base_dataset.py index a004b8d..a9e9c24 100644 --- a/text_recognizer/datasets/base_dataset.py +++ b/text_recognizer/datasets/base_dataset.py @@ -61,13 +61,13 @@ def convert_strings_to_labels( strings: Sequence[str], mapping: Dict[str, int], length: int ) -> Tensor: """ - Convert a sequence of N strings to (N, length) ndarray, with each string wrapped with <S> and </S> tokens, - and padded wiht the <P> token. + Convert a sequence of N strings to (N, length) ndarray, with each string wrapped with <s> and </s> tokens, + and padded wiht the <p> token. """ - labels = torch.ones((len(strings), length), dtype=torch.long) * mapping["<P>"] + labels = torch.ones((len(strings), length), dtype=torch.long) * mapping["<p>"] for i, string in enumerate(strings): tokens = list(string) - tokens = ["<S>", *tokens, "</S>"] + tokens = ["<s>", *tokens, "</s>"] for j, token in enumerate(tokens): labels[i, j] = mapping[token] return labels diff --git a/text_recognizer/datasets/emnist.py b/text_recognizer/datasets/emnist.py index 7c208c4..66101b5 100644 --- a/text_recognizer/datasets/emnist.py +++ b/text_recognizer/datasets/emnist.py @@ -70,9 +70,11 @@ class EMNIST(BaseDataModule): if stage == "fit" or stage is None: with h5py.File(PROCESSED_DATA_FILENAME, "r") as f: self.x_train = f["x_train"][:] - self.y_train = f["y_train"][:] + self.y_train = f["y_train"][:].squeeze().astype(int) - dataset_train = BaseDataset(self.x_train, self.y_train, transform=self.transform) + dataset_train = BaseDataset( + self.x_train, self.y_train, transform=self.transform + ) train_size = int(self.train_fraction * len(dataset_train)) val_size = len(dataset_train) - train_size self.data_train, self.data_val = random_split( @@ -82,8 +84,10 @@ class EMNIST(BaseDataModule): if stage == "test" or stage is None: with h5py.File(PROCESSED_DATA_FILENAME, "r") as f: self.x_test = f["x_test"][:] - self.y_test = f["y_test"][:] - self.data_test = BaseDataset(self.x_test, self.y_test, transform=self.transform) + self.y_test = f["y_test"][:].squeeze().astype(int) + self.data_test = BaseDataset( + self.x_test, self.y_test, transform=self.transform + ) def __repr__(self) -> str: basic = f"EMNIST Dataset\nNum classes: {len(self.mapping)}\nMapping: {self.mapping}\nDims: {self.dims}\n" diff --git a/text_recognizer/datasets/emnist_essentials.json b/text_recognizer/datasets/emnist_essentials.json index 100b36a..3f46a73 100644 --- a/text_recognizer/datasets/emnist_essentials.json +++ b/text_recognizer/datasets/emnist_essentials.json @@ -1 +1 @@ -{"characters": ["<b>", "<s>", "</s>", "<p>", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", " ", "!", "\"", "#", "&", "'", "(", ")", "*", "+", ",", "-", ".", "/", ":", ";", "?"], "input_shape": [28, 28]}
\ No newline at end of file +{"characters": ["<b>", "<s>", "</s>", "<p>", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", " ", "!", "\"", "#", "&", "'", "(", ")", "*", "+", ",", "-", ".", "/", ":", ";", "?"], "input_shape": [28, 28]} diff --git a/text_recognizer/datasets/emnist_lines.py b/text_recognizer/datasets/emnist_lines.py index ae23feb..9ebad22 100644 --- a/text_recognizer/datasets/emnist_lines.py +++ b/text_recognizer/datasets/emnist_lines.py @@ -1,16 +1,21 @@ """Dataset of generated text from EMNIST characters.""" from collections import defaultdict from pathlib import Path -from typing import Dict, Sequence +from typing import Callable, Dict, Tuple, Sequence import h5py from loguru import logger import numpy as np +from PIL import Image import torch from torchvision import transforms +from torchvision.transforms.functional import InterpolationMode -from text_recognizer.datasets.base_dataset import BaseDataset -from text_recognizer.datasets.base_data_module import BaseDataModule +from text_recognizer.datasets.base_dataset import BaseDataset, convert_strings_to_labels +from text_recognizer.datasets.base_data_module import ( + BaseDataModule, + load_and_print_info, +) from text_recognizer.datasets.emnist import EMNIST from text_recognizer.datasets.sentence_generator import SentenceGenerator @@ -54,18 +59,23 @@ class EMNISTLines(BaseDataModule): self.emnist = EMNIST() self.mapping = self.emnist.mapping - max_width = int(self.emnist.dims[2] * (self.max_length + 1) * (1 - self.min_overlap)) + IMAGE_X_PADDING - - if max_width <= IMAGE_WIDTH: - raise ValueError("max_width greater than IMAGE_WIDTH") + max_width = ( + int(self.emnist.dims[2] * (self.max_length + 1) * (1 - self.min_overlap)) + + IMAGE_X_PADDING + ) + + if max_width >= IMAGE_WIDTH: + raise ValueError( + f"max_width {max_width} greater than IMAGE_WIDTH {IMAGE_WIDTH}" + ) self.dims = ( self.emnist.dims[0], - self.emnist.dims[1], - self.emnist.dims[2] * self.max_length, + IMAGE_HEIGHT, + IMAGE_WIDTH ) - if self.max_length <= MAX_OUTPUT_LENGTH: + if self.max_length >= MAX_OUTPUT_LENGTH: raise ValueError("max_length greater than MAX_OUTPUT_LENGTH") self.output_dims = (MAX_OUTPUT_LENGTH, 1) @@ -77,8 +87,11 @@ class EMNISTLines(BaseDataModule): def data_filename(self) -> Path: """Return name of dataset.""" return ( - DATA_DIRNAME - / f"ml_{self.max_length}_o{self.min_overlap:f}_{self.max_overlap:f}_ntr{self.num_train}_ntv{self.num_val}_nte{self.num_test}_{self.with_start_end_tokens}.h5" + DATA_DIRNAME / (f"ml_{self.max_length}_" + f"o{self.min_overlap:f}_{self.max_overlap:f}_" + f"ntr{self.num_train}_" + f"ntv{self.num_val}_" + f"nte{self.num_test}.h5") ) def prepare_data(self) -> None: @@ -92,21 +105,28 @@ class EMNISTLines(BaseDataModule): def setup(self, stage: str = None) -> None: logger.info("EMNISTLinesDataset loading data from HDF5...") if stage == "fit" or stage is None: + print(self.data_filename) with h5py.File(self.data_filename, "r") as f: x_train = f["x_train"][:] y_train = torch.LongTensor(f["y_train"][:]) x_val = f["x_val"][:] y_val = torch.LongTensor(f["y_val"][:]) - self.data_train = BaseDataset(x_train, y_train, transform=_get_transform(augment=self.augment)) - self.data_val = BaseDataset(x_val, y_val, transform=_get_transform(augment=self.augment)) + self.data_train = BaseDataset( + x_train, y_train, transform=_get_transform(augment=self.augment) + ) + self.data_val = BaseDataset( + x_val, y_val, transform=_get_transform(augment=self.augment) + ) if stage == "test" or stage is None: with h5py.File(self.data_filename, "r") as f: x_test = f["x_test"][:] y_test = torch.LongTensor(f["y_test"][:]) - self.data_train = BaseDataset(x_test, y_test, transform=_get_transform(augment=False)) + self.data_test = BaseDataset( + x_test, y_test, transform=_get_transform(augment=False) + ) def __repr__(self) -> str: """Return str about dataset.""" @@ -132,53 +152,129 @@ class EMNISTLines(BaseDataModule): def _generate_data(self, split: str) -> None: logger.info(f"EMNISTLines generating data for {split}...") - sentence_generator = SentenceGenerator(self.max_length - 2) # Subtract by 2 because start/end token + sentence_generator = SentenceGenerator( + self.max_length - 2 + ) # Subtract by 2 because start/end token emnist = self.emnist emnist.prepare_data() emnist.setup() if split == "train": - samples_by_char = _get_samples_by_char(emnist.x_train, emnist.y_train, emnist.mapping) + samples_by_char = _get_samples_by_char( + emnist.x_train, emnist.y_train, emnist.mapping + ) num = self.num_train elif split == "val": - samples_by_char = _get_samples_by_char(emnist.x_train, emnist.y_train, emnist.mapping) + samples_by_char = _get_samples_by_char( + emnist.x_train, emnist.y_train, emnist.mapping + ) num = self.num_val - elif split == "test": - samples_by_char = _get_samples_by_char(emnist.x_test, emnist.y_test, emnist.mapping) + else: + samples_by_char = _get_samples_by_char( + emnist.x_test, emnist.y_test, emnist.mapping + ) num = self.num_test DATA_DIRNAME.mkdir(parents=True, exist_ok=True) - with h5py.File(self.data_filename, "w") as f: + with h5py.File(self.data_filename, "a") as f: x, y = _create_dataset_of_images( - num, samples_by_char, sentence_generator, self.min_overlap, self.max_overlap, self.dims - ) - y = _convert_strings_to_labels( - y, - emnist.inverse_mapping, - length=MAX_OUTPUT_LENGTH - ) + num, + samples_by_char, + sentence_generator, + self.min_overlap, + self.max_overlap, + self.dims, + ) + y = convert_strings_to_labels( + y, emnist.inverse_mapping, length=MAX_OUTPUT_LENGTH + ) f.create_dataset(f"x_{split}", data=x, dtype="u1", compression="lzf") f.create_dataset(f"y_{split}", data=y, dtype="u1", compression="lzf") -def _get_samples_by_char(samples: np.ndarray, labels: np.ndarray, mapping: Dict) -> defaultdict: + +def _get_samples_by_char( + samples: np.ndarray, labels: np.ndarray, mapping: Dict +) -> defaultdict: samples_by_char = defaultdict(list) for sample, label in zip(samples, labels): samples_by_char[mapping[label]].append(sample) return samples_by_char -def _construct_image_from_string(): - pass - - def _select_letter_samples_for_string(string: str, samples_by_char: defaultdict): - pass - - -def _create_dataset_of_images(num_samples: int, samples_by_char: defaultdict, sentence_generator: SentenceGenerator, min_overlap: float, max_overlap: float, dims: Tuple) -> Tuple[torch.Tensor, torch.Tensor]: + null_image = torch.zeros((28, 28), dtype=torch.uint8) + sample_image_by_char = {} + for char in string: + if char in sample_image_by_char: + continue + samples = samples_by_char[char] + sample = samples[np.random.choice(len(samples))] if samples else null_image + sample_image_by_char[char] = sample.reshape(28, 28) + return [sample_image_by_char[char] for char in string] + + +def _construct_image_from_string( + string: str, + samples_by_char: defaultdict, + min_overlap: float, + max_overlap: float, + width: int, +) -> torch.Tensor: + overlap = np.random.uniform(min_overlap, max_overlap) + sampled_images = _select_letter_samples_for_string(string, samples_by_char) + N = len(sampled_images) + H, W = sampled_images[0].shape + next_overlap_width = W - int(overlap * W) + concatenated_image = torch.zeros((H, width), dtype=torch.uint8) + x = IMAGE_X_PADDING + for image in sampled_images: + concatenated_image[:, x : (x + W)] += image + x += next_overlap_width + return torch.minimum(torch.Tensor([255]), concatenated_image) + + +def _create_dataset_of_images( + num_samples: int, + samples_by_char: defaultdict, + sentence_generator: SentenceGenerator, + min_overlap: float, + max_overlap: float, + dims: Tuple, +) -> Tuple[torch.Tensor, torch.Tensor]: images = torch.zeros((num_samples, IMAGE_HEIGHT, dims[2])) labels = [] for n in range(num_samples): label = sentence_generator.generate() - crop = _construct_image_from_string() + crop = _construct_image_from_string( + label, samples_by_char, min_overlap, max_overlap, dims[-1] + ) + height = crop.shape[0] + y = (IMAGE_HEIGHT - height) // 2 + images[n, y : (y + height), :] = crop + labels.append(label) + return images, labels + + +def _get_transform(augment: bool = False) -> Callable: + if not augment: + return transforms.Compose([transforms.ToTensor()]) + return transforms.Compose( + [ + transforms.ToTensor(), + transforms.ColorJitter(brightness=(0.5, 1.0)), + transforms.RandomAffine( + degrees=3, + translate=(0.0, 0.05), + scale=(0.4, 1.1), + shear=(-40, 50), + interpolation=InterpolationMode.BILINEAR, + fill=0, + ), + ] + ) + + +def generate_emnist_lines() -> None: + """Generates a synthetic handwritten dataset and displays info,""" + load_and_print_info(EMNISTLines) diff --git a/text_recognizer/datasets/sentence_generator.py b/text_recognizer/datasets/sentence_generator.py index dd76652..53b781c 100644 --- a/text_recognizer/datasets/sentence_generator.py +++ b/text_recognizer/datasets/sentence_generator.py @@ -11,7 +11,7 @@ import numpy as np from text_recognizer.datasets.util import DATA_DIRNAME -NLTK_DATA_DIRNAME = DATA_DIRNAME / "raw" / "nltk" +NLTK_DATA_DIRNAME = DATA_DIRNAME / "downloaded" / "nltk" class SentenceGenerator: @@ -47,18 +47,22 @@ class SentenceGenerator: raise ValueError( "Must provide max_length to this method or when making this object." ) - - index = np.random.randint(0, len(self.word_start_indices) - 1) - start_index = self.word_start_indices[index] - end_index_candidates = [] - for index in range(index + 1, len(self.word_start_indices)): - if self.word_start_indices[index] - start_index > max_length: - break - end_index_candidates.append(self.word_start_indices[index]) - end_index = np.random.choice(end_index_candidates) - sampled_text = self.corpus[start_index:end_index].strip() - padding = "_" * (max_length - len(sampled_text)) - return sampled_text + padding + + for _ in range(10): + try: + index = np.random.randint(0, len(self.word_start_indices) - 1) + start_index = self.word_start_indices[index] + end_index_candidates = [] + for index in range(index + 1, len(self.word_start_indices)): + if self.word_start_indices[index] - start_index > max_length: + break + end_index_candidates.append(self.word_start_indices[index]) + end_index = np.random.choice(end_index_candidates) + sampled_text = self.corpus[start_index:end_index].strip() + return sampled_text + except Exception: + pass + raise RuntimeError("Was not able to generate a valid string") def brown_corpus() -> str: |