Refactor EMNIST dataset

5 files changed, 336 insertions, 132 deletions

diff --git a/text_recognizer/datasets/base_data_module.py b/text_recognizer/datasets/base_data_module.py
new file mode 100644
index 0000000..09a0a43
--- /dev/null
+++ b/text_recognizer/datasets/base_data_module.py
@@ -0,0 +1,69 @@
+"""Base lightning DataModule class."""
+from pathlib import Path
+from typing import Dict
+
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+
+
+def load_and_print_info(data_module_class: type) -> None:
+    """Load EMNISTLines and prints info."""
+    dataset = data_module_class()
+    dataset.prepare_data()
+    dataset.setup()
+    print(dataset)
+
+
+class BaseDataModule(pl.LightningDataModule):
+    """Base PyTorch Lightning DataModule."""
+    
+    def __init__(self, batch_size: int = 128, num_workers: int = 0) -> None:
+        super().__init__()
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+        # Placeholders for subclasses.
+        self.dims = None
+        self.output_dims = None
+        self.mapping = None
+
+    @classmethod
+    def data_dirname(cls) -> Path:
+        """Return the path to the base data directory."""
+        return Path(__file__).resolve().parents[2] / "data"
+
+    def config(self) -> Dict:
+        """Return important settings of the dataset."""
+        return {"input_dim": self.dims, "output_dims": self.output_dims, "mapping": self.mapping}
+
+    def prepare_data(self) -> None:
+        """Prepare data for training."""
+        pass
+
+    def setup(self, stage: Any = 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.
+
+        Args:
+            stage (Any): Variable to set splits.
+
+        """
+        self.data_train = None
+        self.data_val = None
+        self.data_test = None
+
+
+    def train_dataloader(self) -> DataLoader:
+        """Retun DataLoader for train data."""
+        return DataLoader(self.data_train, shuffle=True, batch_size=self.batch_size, num_workers=self.num_workers, pin_memory=True)
+
+    def val_dataloader(self) -> DataLoader:
+        """Return DataLoader for val data."""
+        return DataLoader(self.data_val, shuffle=False, batch_size=self.batch_size, num_workers=self.num_workers, pin_memory=True)
+
+    def test_dataloader(self) -> DataLoader:
+        """Return DataLoader for val data."""
+        return DataLoader(self.data_test, shuffle=False, batch_size=self.batch_size, num_workers=self.num_workers, pin_memory=True)
+
diff --git a/text_recognizer/datasets/download_utils.py b/text_recognizer/datasets/download_utils.py
new file mode 100644
index 0000000..7a2cab8
--- /dev/null
+++ b/text_recognizer/datasets/download_utils.py
@@ -0,0 +1,73 @@
+"""Util functions for downloading datasets."""
+import hashlib
+from pathlib import Path
+from typing import Dict, List, Optional
+from urllib.request import urlretrieve
+
+from loguru import logger
+from tqdm import tqdm
+
+
+def _compute_sha256(filename: Path) -> str:
+    """Returns the SHA256 checksum of a file."""
+    with filename.open(mode="rb") as f:
+        return hashlib.sha256(f.read()).hexdigest()
+
+
+class TqdmUpTo(tqdm):
+    """TQDM progress bar when downloading files.
+
+    From https://github.com/tqdm/tqdm/blob/master/examples/tqdm_wget.py
+
+    """
+
+    def update_to(
+        self, blocks: int = 1, block_size: int = 1, total_size: Optional[int] = None
+    ) -> None:
+        """Updates the progress bar.
+
+        Args:
+            blocks (int): Number of blocks transferred so far. Defaults to 1.
+            block_size (int): Size of each block, in tqdm units. Defaults to 1.
+            total_size (Optional[int]): Total size in tqdm units. Defaults to None.
+        """
+        if total_size is not None:
+            self.total = total_size  # pylint: disable=attribute-defined-outside-init
+        self.update(blocks * block_size - self.n)
+
+
+def _download_url(url: str, filename: str) -> None:
+    """Downloads a file from url to filename, with a progress bar."""
+    with TqdmUpTo(unit="B", unit_scale=True, unit_divisor=1024, miniters=1) as t:
+        urlretrieve(url, filename, reporthook=t.update_to, data=None)  # nosec
+
+
+def download_dataset(metadata: Dict, dl_dir: Path) -> Optional[Path]:
+    """Downloads dataset using a metadata file.
+
+    Args:
+        metadata (Dict): A metadata file of the dataset.
+        dl_dir (Path): Download directory for the dataset.
+
+    Returns:
+        Optional[Path]: Returns filename if dataset is downloaded, None if it already
+            exists.
+
+    Raises:
+        ValueError: If the SHA-256 value is not the same between the dataset and
+            the metadata file.
+
+    """
+    dl_dir.mkdir(parents=True, exist_ok=True)
+    filename = dl_dir / metadata["filename"]
+    if filename.exists():
+        return
+    logger.info(f"Downloading raw dataset from {metadata['url']} to {filename}...")
+    _download_url(metadata["url"], filename) 
+    logger.info("Computing the SHA-256...")
+    sha256 = _compute_sha256(filename)
+    if sha256 != metadata["sha256"]:
+        raise ValueError(
+                "Downloaded data file SHA-256 does not match that listed in metadata document."
+                )
+    return filename
diff --git a/text_recognizer/datasets/emnist.py b/text_recognizer/datasets/emnist.py
new file mode 100644
index 0000000..e99dbfd
--- /dev/null
+++ b/text_recognizer/datasets/emnist.py
@@ -0,0 +1,194 @@
+"""EMNIST dataset: downloads it from FSDL aws url if not present."""
+from pathlib import Path
+from typing import Sequence, Tuple
+import json
+import os
+import shutil
+import zipfile
+
+import h5py
+import numpy as np
+from loguru import logger
+import toml
+import torch
+from torch.utils.data import random_split
+from torchvision import transforms
+
+from text_recognizer.datasets.base_dataset import BaseDataset
+from text_recognizer.datasets.base_data_module import BaseDataModule, load_print_info
+from text_recognizer.datasets.download_utils import download_dataset
+
+
+SEED = 4711
+NUM_SPECIAL_TOKENS = 4
+SAMPLE_TO_BALANCE = True 
+
+RAW_DATA_DIRNAME = BaseDataModule.data_dirname() / "raw" / "emnist"
+METADATA_FILENAME = RAW_DATA_DIRNAME / "metadata.toml"
+DL_DATA_DIRNAME = BaseDataModule.data_dirname() / "downloaded" / "emnist"
+PROCESSED_DATA_DIRNAME = BaseDataset.data_dirname() / "processed" / "emnist"
+PROCESSED_DATA_FILENAME = PROCESSED_DATA_DIRNAME / "byclass.h5"
+ESSENTIALS_FILENAME = Path(__file__).parents[0].resolve() / "emnsit_essentials.json"
+
+
+class EMNIST(BaseDataModule):
+    """
+    "The EMNIST dataset is a set of handwritten character digits derived from the NIST Special Database 19
+    and converted to a 28x28 pixel image format and dataset structure that directly matches the MNIST dataset."
+    From https://www.nist.gov/itl/iad/image-group/emnist-dataset
+
+    The data split we will use is
+    EMNIST ByClass: 814,255 characters. 62 unbalanced classes.
+    """
+
+    def __init__(self, batch_size: int = 128, num_workers: int = 0, train_fraction: float = 0.8) -> None:
+        super().__init__(batch_size, num_workers)
+        if not ESSENTIALS_FILENAME.exists():
+            _download_and_process_emnist()
+        with ESSENTIALS_FILENAME.open() as f:
+            essentials = json.load(f)
+        self.train_fraction = train_fraction
+        self.mapping = list(essentials["characters"])
+        self.inverse_mapping = {v: k for k, v in enumerate(self.mapping)}
+        self.data_train = None
+        self.data_val = None
+        self.data_test = None
+        self.transform = transforms.Compose([transforms.ToTensor()])
+        self.dims = (1, *essentials["input_shape"])
+        self.output_dims = (1,)
+
+    def prepare_data(self) -> None:
+        if not PROCESSED_DATA_FILENAME.exists():
+            _download_and_process_emnist()
+
+    def setup(self, stage: str = None) -> None:
+        if stage == "fit" or stage is None:
+            with h5py.File(PROCESSED_DATA_FILENAME, "r") as f:
+                data = f["x_train"][:]
+                targets = f["y_train"][:]
+        
+            dataset_train = BaseDataset(data, targets, 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(dataset_train, [train_size, val_size], generator=torch.Generator())
+
+        if stage == "test" or stage is None:
+            with h5py.File(PROCESSED_DATA_FILENAME, "r") as f:
+                data = f["x_test"][:]
+                targets = f["y_test"][:]
+            self.data_test = BaseDataset(data, targets, transform=self.transform)
+
+
+    def __repr__(self) -> str:
+        basic = f"EMNIST Dataset\nNum classes: {len(self.mapping)}\nMapping: {self.mapping}\nDims: {self.dims}\n"
+        if not any([self.data_train, self.data_val, self.data_test]):
+            return basic
+
+        datum, target = next(iter(self.train_dataloader()))
+        data = (
+            f"Train/val/test sizes: {len(self.data_train)}, {len(self.data_val)}, {len(self.data_test)}\n"
+            f"Batch x stats: {(datum.shape, datum.dtype, datum.min(), datum.mean(), datum.std(), datum.max())}\n"
+            f"Batch y stats: {(target.shape, target.dtype, target.min(), target.max())}\n"
+        )
+
+        return basic + data
+
+
+def _download_and_process_emnist() -> None:
+    metadata = toml.load(METADATA_FILENAME)
+    download_dataset(metadata, DL_DATA_DIRNAME)
+    _process_raw_dataset(metadata["filename"], DL_DATA_DIRNAME)
+
+
+def _process_raw_dataset(filename: str, dirname: Path) -> None:
+    logger.info("Unzipping EMNIST...")
+    curdir = os.getcwd()
+    os.chdir(dirname)
+    content = zipfile.ZipFile(filename, "r")
+    content.extract("matlab/emnist-byclass.mat")
+
+    from scipy.io import loadmat
+
+    logger.info("Loading training data from .mat file")
+    data = loadmat("matlab/emnist-byclass.mat")
+    x_train = data["dataset"]["train"][0, 0]["images"][0, 0].reshape(-1, 28, 28).swapaxes(1, 2)
+    y_train = data["dataset"]["train"][0, 0]["labels"][0, 0] + NUM_SPECIAL_TOKENS
+    x_test = data["dataset"]["test"][0, 0]["images"][0, 0].reshape(-1, 28, 28).swapaxes(1, 2)
+    y_test = data["dataset"]["test"][0, 0]["labels"][0, 0] + NUM_SPECIAL_TOKENS
+
+    if SAMPLE_TO_BALANCE:
+        logger.info("Balancing classes to reduce amount of data")
+        x_train, y_train = _sample_to_balance(x_train, y_train)
+        x_test, y_test = _sample_to_balance(x_test, y_test)
+
+
+    logger.info("Saving to HDF5 in a compressed format...")
+    PROCESSED_DATA_DIRNAME.mkdir(parents=True, exist_ok=True)
+    with h5py.File(PROCESSED_DATA_FILENAME, "w") as f:
+        f.create_dataset("x_train", data=x_train, dtype="u1", compression="lzf")
+        f.create_dataset("y_train", data=y_train, dtype="u1", compression="lzf")
+        f.create_dataset("x_test", data=x_test, dtype="u1", compression="lzf")
+        f.create_dataset("y_test", data=y_test, dtype="u1", compression="lzf")
+
+    logger.info("Saving essential dataset parameters to text_recognizer/datasets...")
+    mapping = {int(k): chr(v) for k, v in data["dataset"]["mapping"][0, 0]}
+    characters = _augment_emnist_characters(mapping.values())
+    essentials = {"characters": characters, "input_shape": list(x_train.shape[1:])}
+
+    with ESSENTIALS_FILENAME.open(mode="w") as f:
+        json.dump(essentials, f)
+
+    logger.info("Cleaning up...")
+    shutil.rmtree("matlab")
+    os.chdir(curdir)
+
+
+def _sample_to_balance(x: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Balances the dataset by taking the mean number of instances per class."""
+    np.random.seed(SEED)
+    num_to_sample = int(np.bincount(y.flatten()).mean())
+    all_sampled_indices = []
+    for label in np.unique(y.flatten()):
+        indices = np.where(y == label)[0]
+        sampled_indices = np.unique(np.random.choice(indices, num_to_sample))
+        all_sampled_indices.append(sampled_indices)
+    indices = np.concatenate(all_sampled_indices)
+    x_sampled = x[indices]
+    y_sampled= y[indices]
+    return x_sampled, y_sampled
+
+
+def _augment_emnist_characters(characters: Sequence[str]) -> Sequence[str]:
+    """Augment the mapping with extra symbols."""
+    # Extra characters from the IAM dataset.
+    iam_characters = [
+            " ",
+            "!",
+            '"',
+            "#",
+            "&",
+            "'",
+            "(",
+            ")",
+            "*",
+            "+",
+            ",",
+            "-",
+            ".",
+            "/",
+            ":",
+            ";",
+            "?",
+        ]
+
+    # Also add special tokens for:
+    # - CTC blank token at index 0
+    # - Start token at index 1
+    # - End token at index 2
+    # - Padding token at index 3
+    # Note: Do not forget to update NUM_SPECIAL_TOKENS if changing this!
+    return ["<b>", "<s>", "</s>", "<p>", *characters, *iam_characters]
+
+
+if __name__ == "__main__":
+    load_print_info(EMNIST)
diff --git a/text_recognizer/datasets/emnist_dataset.py b/text_recognizer/datasets/emnist_dataset.py
deleted file mode 100644
index 9884fdf..0000000
--- a/text_recognizer/datasets/emnist_dataset.py
+++ /dev/null
@@ -1,131 +0,0 @@
-"""Emnist dataset: black and white images of handwritten characters (Aa-Zz) and digits (0-9)."""
-
-import json
-from pathlib import Path
-from typing import Callable, Optional, Tuple, Union
-
-from loguru import logger
-import numpy as np
-from PIL import Image
-import torch
-from torch import Tensor
-from torchvision.datasets import EMNIST
-from torchvision.transforms import Compose, ToTensor
-
-from text_recognizer.datasets.dataset import Dataset
-from text_recognizer.datasets.transforms import Transpose
-from text_recognizer.datasets.util import DATA_DIRNAME
-
-
-class EmnistDataset(Dataset):
-    """This is a class for resampling and subsampling the PyTorch EMNIST dataset."""
-
-    def __init__(
-        self,
-        pad_token: str = None,
-        train: bool = False,
-        sample_to_balance: bool = False,
-        subsample_fraction: float = None,
-        transform: Optional[Callable] = None,
-        target_transform: Optional[Callable] = None,
-        seed: int = 4711,
-    ) -> None:
-        """Loads the dataset and the mappings.
-
-        Args:
-            pad_token (str): The pad token symbol. Defaults to _.
-            train (bool): If True, loads the training set, otherwise the validation set is loaded. Defaults to False.
-            sample_to_balance (bool): Resamples the dataset to make it balanced. Defaults to False.
-            subsample_fraction (float): Description of parameter `subsample_fraction`. Defaults to None.
-            transform (Optional[Callable]): Transform(s) for input data. Defaults to None.
-            target_transform (Optional[Callable]): Transform(s) for output data. Defaults to None.
-            seed (int): Seed number. Defaults to 4711.
-
-        """
-        super().__init__(
-            train=train,
-            subsample_fraction=subsample_fraction,
-            transform=transform,
-            target_transform=target_transform,
-            pad_token=pad_token,
-        )
-
-        self.sample_to_balance = sample_to_balance
-
-        # Have to transpose the emnist characters, ToTensor norms input between [0,1].
-        if transform is None:
-            self.transform = Compose([Transpose(), ToTensor()])
-
-        self.target_transform = None
-
-        self.seed = seed
-
-    def __getitem__(self, index: Union[int, Tensor]) -> Tuple[Tensor, Tensor]:
-        """Fetches samples from the dataset.
-
-        Args:
-            index (Union[int, Tensor]): The indices of the samples to fetch.
-
-        Returns:
-            Tuple[Tensor, Tensor]: Data target tuple.
-
-        """
-        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 Dataset\n"
-            f"Num classes: {self.num_classes}\n"
-            f"Input shape: {self.input_shape}\n"
-            f"Mapping: {self.mapper.mapping}\n"
-        )
-
-    def _sample_to_balance(self) -> None:
-        """Because the dataset is not balanced, we take at most the mean number of instances per class."""
-        np.random.seed(self.seed)
-        x = self._data
-        y = self._targets
-        num_to_sample = int(np.bincount(y.flatten()).mean())
-        all_sampled_indices = []
-        for label in np.unique(y.flatten()):
-            inds = np.where(y == label)[0]
-            sampled_indices = np.unique(np.random.choice(inds, num_to_sample))
-            all_sampled_indices.append(sampled_indices)
-        indices = np.concatenate(all_sampled_indices)
-        x_sampled = x[indices]
-        y_sampled = y[indices]
-        self._data = x_sampled
-        self._targets = y_sampled
-
-    def load_or_generate_data(self) -> None:
-        """Fetch the EMNIST dataset."""
-        dataset = EMNIST(
-            root=DATA_DIRNAME,
-            split="byclass",
-            train=self.train,
-            download=False,
-            transform=None,
-            target_transform=None,
-        )
-
-        self._data = dataset.data
-        self._targets = dataset.targets
-
-        if self.sample_to_balance:
-            self._sample_to_balance()
-
-        if self.subsample_fraction is not None:
-            self._subsample()
diff --git a/text_recognizer/datasets/emnist_essentials.json b/text_recognizer/datasets/emnist_essentials.json
deleted file mode 100644
index 2a0648a..0000000
--- a/text_recognizer/datasets/emnist_essentials.json
+++ /dev/null
@@ -1 +0,0 @@
-{"mapping": [[0, "0"], [1, "1"], [2, "2"], [3, "3"], [4, "4"], [5, "5"], [6, "6"], [7, "7"], [8, "8"], [9, "9"], [10, "A"], [11, "B"], [12, "C"], [13, "D"], [14, "E"], [15, "F"], [16, "G"], [17, "H"], [18, "I"], [19, "J"], [20, "K"], [21, "L"], [22, "M"], [23, "N"], [24, "O"], [25, "P"], [26, "Q"], [27, "R"], [28, "S"], [29, "T"], [30, "U"], [31, "V"], [32, "W"], [33, "X"], [34, "Y"], [35, "Z"], [36, "a"], [37, "b"], [38, "c"], [39, "d"], [40, "e"], [41, "f"], [42, "g"], [43, "h"], [44, "i"], [45, "j"], [46, "k"], [47, "l"], [48, "m"], [49, "n"], [50, "o"], [51, "p"], [52, "q"], [53, "r"], [54, "s"], [55, "t"], [56, "u"], [57, "v"], [58, "w"], [59, "x"], [60, "y"], [61, "z"]], "input_shape": [28, 28]}