path: root/text_recognizer/datasets/emnist_lines_dataset.py

"""Emnist Lines dataset: synthetic handwritten lines dataset made from Emnist characters."""

from collections import defaultdict
from pathlib import Path
from typing import Callable, Dict, List, Optional, Tuple, Union

import click
import h5py
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
from text_recognizer.datasets.emnist_dataset import EmnistDataset, Transpose
from text_recognizer.datasets.sentence_generator import SentenceGenerator
from text_recognizer.datasets.util import (
    DATA_DIRNAME,
    EmnistMapper,
    ESSENTIALS_FILENAME,
)

DATA_DIRNAME = DATA_DIRNAME / "processed" / "emnist_lines"

MAX_WIDTH = 952


class EmnistLinesDataset(Dataset):
    """Synthetic dataset of lines from the Brown corpus with Emnist characters."""

    def __init__(
        self,
        train: bool = False,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        subsample_fraction: float = None,
        max_length: int = 34,
        min_overlap: float = 0,
        max_overlap: float = 0.33,
        num_samples: int = 10000,
        seed: int = 4711,
        init_token: Optional[str] = None,
        pad_token: Optional[str] = None,
        eos_token: Optional[str] = None,
        lower: bool = False,
    ) -> None:
        """Set attributes and loads the dataset.

        Args:
            train (bool): Flag for the filename. Defaults to False. Defaults to None.
            transform (Optional[Callable]): The transform of the data. Defaults to None.
            target_transform (Optional[Callable]): The transform of the target. Defaults to None.
            subsample_fraction (float): The fraction of the dataset to use for training. Defaults to None.
            max_length (int): The maximum number of characters. Defaults to 34.
            min_overlap (float): The minimum overlap between concatenated images. Defaults to 0.
            max_overlap (float): The maximum overlap between concatenated images. Defaults to 0.33.
            num_samples (int): Number of samples to generate. Defaults to 10000.
            seed (int): Seed number. Defaults to 4711.
            init_token (Optional[str]): String representing the start of sequence token. Defaults to None.
            pad_token (Optional[str]): String representing the pad token. Defaults to None.
            eos_token (Optional[str]): String representing the end of sequence token. Defaults to None.
            lower (bool): If True, convert uppercase letters to lowercase. Otherwise, use both upper and lowercase.

        """
        self.pad_token = "_" if pad_token is None else pad_token

        super().__init__(
            train=train,
            transform=transform,
            target_transform=target_transform,
            subsample_fraction=subsample_fraction,
            init_token=init_token,
            pad_token=self.pad_token,
            eos_token=eos_token,
            lower=lower,
        )

        # Extract dataset information.
        self._input_shape = self._mapper.input_shape
        self.num_classes = self._mapper.num_classes

        self.max_length = max_length
        self.min_overlap = min_overlap
        self.max_overlap = max_overlap
        self.num_samples = num_samples
        self._input_shape = (
            self.input_shape[0],
            self.input_shape[1] * self.max_length,
        )
        self._output_shape = (self.max_length, self.num_classes)
        self.seed = seed

        # Placeholders for the dataset.
        self._data = None
        self._target = None

    def __getitem__(self, index: Union[int, Tensor]) -> Tuple[Tensor, Tensor]:
        """Fetches data, target pair of the dataset for a given and index or indices.

        Args:
            index (Union[int, Tensor]): Either a list or int of indices/index.

        Returns:
            Tuple[Tensor, Tensor]: Data target pair.

        """
        if torch.is_tensor(index):
            index = index.tolist()

        data = self.data[index]
        targets = self.targets[index]

        if self.transform:
            data = self.transform(data)

        if self.target_transform:
            targets = self.target_transform(targets)

        return data, targets

    def __repr__(self) -> str:
        """Returns information about the dataset."""
        return (
            "EMNIST Lines Dataset\n"  # pylint: disable=no-member
            f"Max length: {self.max_length}\n"
            f"Min overlap: {self.min_overlap}\n"
            f"Max overlap: {self.max_overlap}\n"
            f"Num classes: {self.num_classes}\n"
            f"Input shape: {self.input_shape}\n"
            f"Data: {self.data.shape}\n"
            f"Tagets: {self.targets.shape}\n"
        )

    @property
    def data_filename(self) -> Path:
        """Path to the h5 file."""
        filename = "train.pt" if self.train else "test.pt"
        return DATA_DIRNAME / filename

    def load_or_generate_data(self) -> None:
        """Loads the dataset, if it does not exist a new dataset is generated before loading it."""
        np.random.seed(self.seed)

        if not self.data_filename.exists():
            self._generate_data()
        self._load_data()
        self._subsample()

    def _load_data(self) -> None:
        """Loads the dataset from the h5 file."""
        logger.debug("EmnistLinesDataset loading data from HDF5...")
        with h5py.File(self.data_filename, "r") as f:
            self._data = f["data"][()]
            self._targets = f["targets"][()]

    def _generate_data(self) -> str:
        """Generates a dataset with the Brown corpus and Emnist characters."""
        logger.debug("Generating data...")

        sentence_generator = SentenceGenerator(self.max_length)

        # Load emnist dataset.
        emnist = EmnistDataset(
            train=self.train, sample_to_balance=True, pad_token=self.pad_token
        )
        emnist.load_or_generate_data()

        samples_by_character = get_samples_by_character(
            emnist.data.numpy(), emnist.targets.numpy(), self.mapper.mapping,
        )

        DATA_DIRNAME.mkdir(parents=True, exist_ok=True)
        with h5py.File(self.data_filename, "a") as f:
            data, targets = create_dataset_of_images(
                self.num_samples,
                samples_by_character,
                sentence_generator,
                self.min_overlap,
                self.max_overlap,
            )

            targets = convert_strings_to_categorical_labels(
                targets, emnist.inverse_mapping
            )

            f.create_dataset("data", data=data, dtype="u1", compression="lzf")
            f.create_dataset("targets", data=targets, dtype="u1", compression="lzf")


def get_samples_by_character(
    samples: np.ndarray, labels: np.ndarray, mapping: Dict
) -> defaultdict:
    """Creates a dictionary with character as key and value as the list of images of that character.

    Args:
        samples (np.ndarray): Dataset of images of characters.
        labels (np.ndarray): The labels for each image.
        mapping (Dict): The Emnist mapping dictionary.

    Returns:
        defaultdict: A dictionary with characters as keys and list of images as values.

    """
    samples_by_character = defaultdict(list)
    for sample, label in zip(samples, labels.flatten()):
        samples_by_character[mapping[label]].append(sample)
    return samples_by_character


def select_letter_samples_for_string(
    string: str, samples_by_character: Dict
) -> List[np.ndarray]:
    """Randomly selects Emnist characters to use for the senetence.

    Args:
        string (str): The word or sentence.
        samples_by_character (Dict): The dictionary of emnist images of each character.

    Returns:
        List[np.ndarray]: A list of emnist images of the string.

    """
    zero_image = np.zeros((28, 28), np.uint8)
    sample_image_by_character = {}
    for character in string:
        if character in sample_image_by_character:
            continue
        samples = samples_by_character[character]
        sample = samples[np.random.choice(len(samples))] if samples else zero_image
        sample_image_by_character[character] = sample.reshape(28, 28).swapaxes(0, 1)
    return [sample_image_by_character[character] for character in string]


def construct_image_from_string(
    string: str, samples_by_character: Dict, min_overlap: float, max_overlap: float
) -> np.ndarray:
    """Concatenates images of the characters in the string.

    The concatination is made with randomly selected overlap so that some portion of the character will overlap.

    Args:
        string (str): The word or sentence.
        samples_by_character (Dict): The dictionary of emnist images of each character.
        min_overlap (float): Minimum amount of overlap between Emnist images.
        max_overlap (float): Maximum amount of overlap between Emnist images.

    Returns:
        np.ndarray: The Emnist image of the string.

    """
    overlap = np.random.uniform(min_overlap, max_overlap)
    sampled_images = select_letter_samples_for_string(string, samples_by_character)
    length = len(sampled_images)
    height, width = sampled_images[0].shape
    next_overlap_width = width - int(overlap * width)
    concatenated_image = np.zeros((height, width * length), np.uint8)
    x = 0
    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)


def create_dataset_of_images(
    length: int,
    samples_by_character: Dict,
    sentence_generator: SentenceGenerator,
    min_overlap: float,
    max_overlap: float,
) -> Tuple[np.ndarray, List[str]]:
    """Creates a dataset with images and labels from strings generated from the SentenceGenerator.

    Args:
        length (int): The number of characters for each string.
        samples_by_character (Dict): The dictionary of emnist images of each character.
        sentence_generator (SentenceGenerator): A SentenceGenerator objest.
        min_overlap (float): Minimum amount of overlap between Emnist images.
        max_overlap (float): Maximum amount of overlap between Emnist images.

    Returns:
        Tuple[np.ndarray, List[str]]: A list of Emnist images and a list of the strings (labels).

    Raises:
        RuntimeError: If the sentence generator is not able to generate a string.

    """
    sample_label = sentence_generator.generate()
    sample_image = construct_image_from_string(sample_label, samples_by_character, 0, 0)
    images = np.zeros((length, sample_image.shape[0], sample_image.shape[1]), np.uint8)
    labels = []
    for n in range(length):
        label = None
        # Try several times to generate before actually throwing an error.
        for _ in range(10):
            try:
                label = sentence_generator.generate()
                break
            except Exception:  # pylint: disable=broad-except
                pass
        if label is None:
            raise RuntimeError("Was not able to generate a valid string.")
        images[n] = construct_image_from_string(
            label, samples_by_character, min_overlap, max_overlap
        )
        labels.append(label)
    return images, labels


def convert_strings_to_categorical_labels(
    labels: List[str], mapping: Dict
) -> np.ndarray:
    """Translates a string of characters in to a target array of class int."""
    return np.array([[mapping[c] for c in label] for label in labels])


@click.command()
@click.option(
    "--max_length", type=int, default=34, help="Number of characters in a sentence."
)
@click.option(
    "--min_overlap", type=float, default=0.0, help="Min overlap between characters."
)
@click.option(
    "--max_overlap", type=float, default=0.33, help="Max overlap between characters."
)
@click.option("--num_train", type=int, default=10_000, help="Number of train examples.")
@click.option("--num_test", type=int, default=1_000, help="Number of test examples.")
def create_datasets(
    max_length: int = 34,
    min_overlap: float = 0,
    max_overlap: float = 0.33,
    num_train: int = 10000,
    num_test: int = 1000,
) -> None:
    """Creates a training an validation dataset of Emnist lines."""
    num_samples = [num_train, num_test]
    for num, train in zip(num_samples, [True, False]):
        emnist_lines = EmnistLinesDataset(
            train=train,
            max_length=max_length,
            min_overlap=min_overlap,
            max_overlap=max_overlap,
            num_samples=num,
        )
        emnist_lines.load_or_generate_data()


if __name__ == "__main__":
    create_datasets()