Add IAM paragraphs dataset

1 files changed, 310 insertions, 0 deletions

diff --git a/text_recognizer/data/iam_paragraphs.py b/text_recognizer/data/iam_paragraphs.py
new file mode 100644
index 0000000..402a8d4
--- /dev/null
+++ b/text_recognizer/data/iam_paragraphs.py
@@ -0,0 +1,310 @@
+"""IAM Paragraphs Dataset class."""
+import json
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+from loguru import logger
+import numpy as np
+from PIL import Image, ImageFile, ImageOps
+import torch
+import torchvision.transforms as transforms
+from torchvision.transforms.functional import InterpolationMode
+from tqdm import tqdm
+
+from text_recognizer.data.base_dataset import (
+    BaseDataset,
+    convert_strings_to_labels,
+    split_dataset,
+)
+from text_recognizer.data.base_data_module import BaseDataModule, load_and_print_info
+from text_recognizer.data.emnist import emnist_mapping
+from text_recognizer.data.iam import IAM
+
+
+PROCESSED_DATA_DIRNAME = BaseDataModule.data_dirname() / "processed" / "iam_paragraphs"
+
+NEW_LINE_TOKEN = "\n"
+
+SEED = 4711
+IMAGE_SCALE_FACTOR = 2
+IMAGE_HEIGHT = 1152 // IMAGE_SCALE_FACTOR
+IMAGE_WIDTH = 1280 // IMAGE_SCALE_FACTOR
+MAX_LABEL_LENGTH = 682
+
+
+class IAMParagraphs(BaseDataModule):
+    """IAM handwriting database paragraphs."""
+
+    def __init__(
+        self,
+        batch_size: int = 128,
+        num_workers: int = 0,
+        train_fraction: float = 0.8,
+        augment: bool = True,
+    ) -> None:
+        super().__init__(batch_size, num_workers)
+        # TODO: pass in transform and target transform
+        # TODO: pass in mapping
+        self.augment = augment
+        self.mapping, self.inverse_mapping, _ = emnist_mapping(
+            extra_symbols=[NEW_LINE_TOKEN]
+        )
+        self.train_fraction = train_fraction
+
+        self.dims = (1, IMAGE_HEIGHT, IMAGE_WIDTH)
+        self.output_dims = (MAX_LABEL_LENGTH, 1)
+        self.data_train: BaseDataset = None
+        self.data_val: BaseDataset = None
+        self.data_test: BaseDataset = None
+
+    def prepare_data(self) -> None:
+        """Create data for training/testing."""
+        if PROCESSED_DATA_DIRNAME.exists():
+            return
+
+        logger.info(
+            "Cropping IAM paragraph regions and saving them along with labels..."
+        )
+
+        iam = IAM()
+        iam.prepare_data()
+
+        properties = {}
+        for split in ["train", "test"]:
+            crops, labels = _get_paragraph_crops_and_labels(iam=iam, split=split)
+            _save_crops_and_labels(crops=crops, labels=labels, split=split)
+
+            properties.update(
+                {
+                    id_: {
+                        "crop_shape": crops[id_].size[::-1],
+                        "label_length": len(label),
+                        "num_lines": _num_lines(label),
+                    }
+                    for id_, label in labels.items()
+                }
+            )
+
+        with (PROCESSED_DATA_DIRNAME / "_properties.json").open("w") as f:
+            json.dump(properties, f, indent=4)
+
+    def setup(self, stage: str = None) -> None:
+        """Loads the data for training/testing."""
+
+        def _load_dataset(split: str, augment: bool) -> BaseDataset:
+            crops, labels = _load_processed_crops_and_labels(split)
+            data = [_resize_image(crop, IMAGE_SCALE_FACTOR) for crop in crops]
+            targets = convert_strings_to_labels(
+                strings=labels, mapping=self.inverse_mapping, length=self.output_dims[0]
+            )
+            return BaseDataset(
+                data,
+                targets,
+                transform=_get_transform(image_shape=self.dims[1:], augment=augment),
+            )
+
+        logger.info(f"Loading IAM paragraph regions and lines for {stage}...")
+        _validate_data_dims(input_dims=self.dims, output_dims=self.output_dims)
+
+        if stage == "fit" or stage is None:
+            data_train = _load_dataset(split="train", augment=self.augment)
+            self.data_train, self.data_val = split_dataset(
+                dataset=data_train, fraction=self.train_fraction, seed=SEED
+            )
+
+        if stage == "test" or stage is None:
+            self.data_test = _load_dataset(split="test", augment=False)
+
+    def __repr__(self) -> str:
+        """Return information about the dataset."""
+        basic = (
+            "IAM Paragraphs Dataset\n"
+            f"Num classes: {len(self.mapping)}\n"
+            f"Input dims: {self.dims}\n"
+            f"Output dims: {self.output_dims}\n"
+        )
+
+        if not any([self.data_train, self.data_val, self.data_test]):
+            return basic
+
+        x, y = next(iter(self.train_dataloader()))
+        xt, yt = next(iter(self.test_dataloader()))
+        data = (
+            "Train/val/test sizes: "
+            f"{len(self.data_train)}, "
+            f"{len(self.data_val)}, "
+            f"{len(self.data_test)}\n"
+            f"Train Batch x stats: {(x.shape, x.dtype, x.min(), x.mean(), x.std(), x.max())}\n"
+            f"Train Batch y stats: {(y.shape, y.dtype, y.min(), y.max())}\n"
+            f"Test Batch x stats: {(xt.shape, xt.dtype, xt.min(), xt.mean(), xt.std(), xt.max())}\n"
+            f"Test Batch y stats: {(yt.shape, yt.dtype, yt.min(), yt.max())}\n"
+        )
+        return basic + data
+
+
+def _get_dataset_properties() -> Dict:
+    """Return properties describing the overall dataset."""
+    with (PROCESSED_DATA_DIRNAME / "_properties.json").open("r") as f:
+        properties = json.load(f)
+
+    def _get_property_values(key: str) -> List:
+        return [value[key] for value in properties.values()]
+
+    crop_shapes = np.array(_get_property_values("crop_shape"))
+    aspect_ratio = crop_shapes[:, 1] / crop_shapes[:, 0]
+    return {
+        "label_length": {
+            "min": min(_get_property_values("label_length")),
+            "max": max(_get_property_values("label_length")),
+        },
+        "num_lines": {
+            "min": min(_get_property_values("num_lines")),
+            "max": max(_get_property_values("num_lines")),
+        },
+        "crop_shape": {"min": crop_shapes.min(axis=0), "max": crop_shapes.max(axis=0),},
+        "aspect_ratio": {
+            "min": aspect_ratio.min(axis=0),
+            "max": aspect_ratio.max(axis=0),
+        },
+    }
+
+
+def _validate_data_dims(
+    input_dims: Optional[Tuple[int, ...]], output_dims: Optional[Tuple[int, ...]]
+) -> None:
+    """Validates input and output dimensions against the properties of the dataset."""
+    properties = _get_dataset_properties()
+
+    max_image_shape = properties["crop_shape"]["max"] / IMAGE_SCALE_FACTOR
+    if (
+        input_dims is not None
+        and input_dims[1] < max_image_shape[0]
+        and input_dims[2] < max_image_shape[1]
+    ):
+        raise ValueError(f"{input_dims} less than {max_image_shape}")
+
+    if (
+        output_dims is not None
+        and output_dims[0] < properties["label_length"]["max"] + 2
+    ):
+        raise ValueError(
+            f"{output_dims} less than {properties['label_length']['max'] + 2}"
+        )
+
+
+def _resize_image(image: Image.Image, scale_factor: int) -> Image.Image:
+    """Resize image by scale factor."""
+    if scale_factor == 1:
+        return image
+    return image.resize(
+        (image.width // scale_factor, image.height // scale_factor),
+        resample=Image.BILINEAR,
+    )
+
+
+def _get_paragraph_crops_and_labels(
+    iam: IAM, split: str
+) -> Tuple[Dict[str, Image.Image], Dict[str, str]]:
+    """Load IAM paragraph crops and labels for a given set."""
+    crops = {}
+    labels = {}
+    for form_filename in tqdm(
+        iam.form_filenames, desc=f"Processing {split} paragraphs"
+    ):
+        id_ = form_filename.stem
+        if not iam.split_by_id[id_] == split:
+            continue
+        image = Image.open(form_filename)
+        image = ImageOps.grayscale(image)
+        image = ImageOps.invert(image)
+
+        line_regions = iam.line_regions_by_id[id_]
+        parameter_box = [
+            min([region["x1"] for region in line_regions]),
+            min([region["y1"] for region in line_regions]),
+            max([region["x2"] for region in line_regions]),
+            max([region["y2"] for region in line_regions]),
+        ]
+        lines = iam.line_strings_by_id[id_]
+
+        crops[id_] = image.crop(parameter_box)
+        labels[id_] = NEW_LINE_TOKEN.join(lines)
+
+    if len(crops) != len(labels):
+        raise ValueError(f"Crops ({len(crops)}) does not match labels ({len(labels)})")
+
+    return crops, labels
+
+
+def _save_crops_and_labels(
+    crops: Dict[str, Image.Image], labels: Dict[str, str], split: str
+) -> None:
+    """Save crops, labels, and shapes of crops of a split."""
+    (PROCESSED_DATA_DIRNAME / split).mkdir(parents=True, exist_ok=True)
+
+    with _labels_filename(split).open("w") as f:
+        json.dump(labels, f, indent=4)
+
+    for id_, crop in crops.items():
+        crop.save(_crop_filename(id_, split))
+
+
+def _load_processed_crops_and_labels(
+    split: str,
+) -> Tuple[Sequence[Image.Image], Sequence[str]]:
+    """Load processed crops and labels for given split."""
+    with _labels_filename(split).open("r") as f:
+        labels = json.load(f)
+
+    sorted_ids = sorted(labels.keys())
+    ordered_crops = [
+        Image.open(_crop_filename(id_, split)).convert("L") for id_ in sorted_ids
+    ]
+    ordered_labels = [labels[id_] for id_ in sorted_ids]
+
+    if len(ordered_crops) != len(ordered_labels):
+        raise ValueError(
+            f"Crops ({len(ordered_crops)}) does not match labels ({len(ordered_labels)})"
+        )
+    return ordered_crops, ordered_labels
+
+
+def _get_transform(image_shape: Tuple[int, int], augment: bool) -> transforms.Compose:
+    """Get transformations for images."""
+    if augment:
+        transforms_list = [
+            transforms.RandomCrop(
+                size=image_shape,
+                padding=None,
+                pad_if_needed=True,
+                fill=0,
+                padding_mode="constant",
+            ),
+            transforms.ColorJitter(brightness=(0.8, 1.6)),
+            transforms.RandomAffine(
+                degrees=1, shear=(-10, 10), interpolation=InterpolationMode.BILINEAR,
+            ),
+        ]
+    else:
+        transforms_list = [transforms.CenterCrop(image_shape)]
+    transforms_list.append(transforms.ToTensor())
+    return transforms.Compose(transforms_list)
+
+
+def _labels_filename(split: str) -> Path:
+    """Return filename of processed labels."""
+    return PROCESSED_DATA_DIRNAME / split / "_labels.json"
+
+
+def _crop_filename(id: str, split: str) -> Path:
+    """Return filename of processed crop."""
+    return PROCESSED_DATA_DIRNAME / split / f"{id}.png"
+
+
+def _num_lines(label: str) -> int:
+    """Return the number of lines of text in label."""
+    return label.count("\n") + 1
+
+
+def create_iam_paragraphs() -> None:
+    load_and_print_info(IAMParagraphs)