Segmentation working!

20 files changed, 551 insertions, 361 deletions

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)
 
 
diff --git a/src/text_recognizer/datasets/iam_paragraphs_dataset.py b/src/text_recognizer/datasets/iam_paragraphs_dataset.py
index c1e8fe2..8ba5142 100644
--- a/src/text_recognizer/datasets/iam_paragraphs_dataset.py
+++ b/src/text_recognizer/datasets/iam_paragraphs_dataset.py
@@ -1,4 +1,5 @@
 """IamParagraphsDataset class and functions for data processing."""
+import random
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import click
@@ -71,13 +72,18 @@ class IamParagraphsDataset(Dataset):
         data = self.data[index]
         targets = self.targets[index]
 
+        seed = np.random.randint(SEED)
+        random.seed(seed)  # apply this seed to target tranfsorms
+        torch.manual_seed(seed)  # needed for torchvision 0.7
         if self.transform:
             data = self.transform(data)
 
+        random.seed(seed)  # apply this seed to target tranfsorms
+        torch.manual_seed(seed)  # needed for torchvision 0.7
         if self.target_transform:
             targets = self.target_transform(targets)
 
-        return data, targets
+        return data, targets.long()
 
     @property
     def ids(self) -> Tensor:
diff --git a/src/text_recognizer/datasets/transforms.py b/src/text_recognizer/datasets/transforms.py
index 1ec23dc..016ec80 100644
--- a/src/text_recognizer/datasets/transforms.py
+++ b/src/text_recognizer/datasets/transforms.py
@@ -4,7 +4,7 @@ from PIL import Image
 import torch
 from torch import Tensor
 import torch.nn.functional as F
-from torchvision.transforms import Compose, RandomAffine, ToTensor
+from torchvision.transforms import Compose, RandomAffine, RandomHorizontalFlip, ToTensor
 
 from text_recognizer.datasets.util import EmnistMapper
 
@@ -77,3 +77,19 @@ class ApplyContrast:
         """Apply mask binary mask to input tensor."""
         mask = x > np.random.RandomState().uniform(low=self.low, high=self.high)
         return x * mask
+
+
+class Unsqueeze:
+    """Add a dimension to the tensor."""
+
+    def __call__(self, x: Tensor) -> Tensor:
+        """Adds dim."""
+        return x.unsqueeze(0)
+
+
+class Squeeze:
+    """Removes the first dimension of a tensor."""
+
+    def __call__(self, x: Tensor) -> Tensor:
+        """Removes first dim."""
+        return x.squeeze(0)
diff --git a/src/text_recognizer/models/__init__.py b/src/text_recognizer/models/__init__.py
index bf89404..a645cec 100644
--- a/src/text_recognizer/models/__init__.py
+++ b/src/text_recognizer/models/__init__.py
@@ -2,11 +2,13 @@
 from .base import Model
 from .character_model import CharacterModel
 from .crnn_model import CRNNModel
+from .segmentation_model import SegmentationModel
 from .transformer_model import TransformerModel
 
 __all__ = [
     "CharacterModel",
     "CRNNModel",
     "Model",
+    "SegmentationModel",
     "TransformerModel",
 ]
diff --git a/src/text_recognizer/models/base.py b/src/text_recognizer/models/base.py
index d394b4c..f2cd4b8 100644
--- a/src/text_recognizer/models/base.py
+++ b/src/text_recognizer/models/base.py
@@ -159,7 +159,7 @@ class Model(ABC):
             self.test_dataset = self.dataset(train=False, **self.dataset_args["args"])
             self.test_dataset.load_or_generate_data()
 
-            # Set the flag to true to disable ability to load data agian.
+            # Set the flag to true to disable ability to load data again.
             self.data_prepared = True
 
     def train_dataloader(self) -> DataLoader:
@@ -260,7 +260,7 @@ class Model(ABC):
     @property
     def mapping(self) -> Dict:
         """Returns the mapping between network output and Emnist character."""
-        return self._mapper.mapping
+        return self._mapper.mapping if self._mapper is not None else None
 
     def eval(self) -> None:
         """Sets the network to evaluation mode."""
@@ -341,7 +341,7 @@ class Model(ABC):
         if input_shape is not None:
             summary(self.network, input_shape, depth=depth, device=device)
         elif self._input_shape is not None:
-            input_shape = (1,) + tuple(self._input_shape)
+            input_shape = tuple(self._input_shape)
             summary(self.network, input_shape, depth=depth, device=device)
         else:
             logger.warning("Could not print summary as input shape is not set.")
@@ -427,7 +427,7 @@ class Model(ABC):
             )
             shutil.copyfile(filepath, str(checkpoint_path / "best.pt"))
 
-    def load_weights(self, network_fn: Type[nn.Module]) -> None:
+    def load_weights(self, network_fn: Optional[Type[nn.Module]] = None) -> None:
         """Load the network weights."""
         logger.debug("Loading network with pretrained weights.")
         filename = glob(self.weights_filename)[0]
@@ -441,7 +441,8 @@ class Model(ABC):
         weights = state_dict["model_state"]
 
         # Initializes the network with trained weights.
-        self._network = network_fn(**self._network_args)
+        if network_fn is not None:
+            self._network = network_fn(**self._network_args)
         self._network.load_state_dict(weights)
 
         if "swa_network" in state_dict:
diff --git a/src/text_recognizer/models/segmentation_model.py b/src/text_recognizer/models/segmentation_model.py
new file mode 100644
index 0000000..613108a
--- /dev/null
+++ b/src/text_recognizer/models/segmentation_model.py
@@ -0,0 +1,75 @@
+"""Segmentation model for detecting and segmenting lines."""
+from typing import Callable, Dict, Optional, Type, Union
+
+import numpy as np
+import torch
+from torch import nn
+from torch import Tensor
+from torch.utils.data import Dataset
+from torchvision.transforms import ToTensor
+
+from text_recognizer.models.base import Model
+
+
+class SegmentationModel(Model):
+    """Model for segmenting lines in an image."""
+
+    def __init__(
+        self,
+        network_fn: str,
+        dataset: str,
+        network_args: Optional[Dict] = None,
+        dataset_args: Optional[Dict] = None,
+        metrics: Optional[Dict] = None,
+        criterion: Optional[Callable] = None,
+        criterion_args: Optional[Dict] = None,
+        optimizer: Optional[Callable] = None,
+        optimizer_args: Optional[Dict] = None,
+        lr_scheduler: Optional[Callable] = None,
+        lr_scheduler_args: Optional[Dict] = None,
+        swa_args: Optional[Dict] = None,
+        device: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            network_fn,
+            dataset,
+            network_args,
+            dataset_args,
+            metrics,
+            criterion,
+            criterion_args,
+            optimizer,
+            optimizer_args,
+            lr_scheduler,
+            lr_scheduler_args,
+            swa_args,
+            device,
+        )
+        self.tensor_transform = ToTensor()
+        self.softmax = nn.Softmax(dim=2)
+
+    @torch.no_grad()
+    def predict_on_image(self, image: Union[np.ndarray, Tensor]) -> Tensor:
+        """Predict on a single input."""
+        self.eval()
+
+        if image.dtype is np.uint8:
+            # Converts an image with range [0, 255] with to PyTorch Tensor with range [0, 1].
+            image = self.tensor_transform(image)
+
+        # Rescale image between 0 and 1.
+        if image.dtype is torch.uint8 or image.dtype is torch.int64:
+            # If the image is an unscaled tensor.
+            image = image.type("torch.FloatTensor") / 255
+
+        if not torch.is_tensor(image):
+            image = Tensor(image)
+
+        # Put the image tensor on the device the model weights are on.
+        image = image.to(self.device)
+
+        logits = self.forward(image)
+
+        segmentation_mask = torch.argmax(logits, dim=1)
+
+        return segmentation_mask
diff --git a/src/text_recognizer/models/transformer_model.py b/src/text_recognizer/models/transformer_model.py
index 968a047..a912122 100644
--- a/src/text_recognizer/models/transformer_model.py
+++ b/src/text_recognizer/models/transformer_model.py
@@ -18,8 +18,8 @@ class TransformerModel(Model):
 
     def __init__(
         self,
-        network_fn: Type[nn.Module],
-        dataset: Type[Dataset],
+        network_fn: str,
+        dataset: str,
         network_args: Optional[Dict] = None,
         dataset_args: Optional[Dict] = None,
         metrics: Optional[Dict] = None,
diff --git a/src/text_recognizer/networks/__init__.py b/src/text_recognizer/networks/__init__.py
index 1635039..f958672 100644
--- a/src/text_recognizer/networks/__init__.py
+++ b/src/text_recognizer/networks/__init__.py
@@ -3,11 +3,13 @@ from .cnn_transformer import CNNTransformer
 from .crnn import ConvolutionalRecurrentNetwork
 from .ctc import greedy_decoder
 from .densenet import DenseNet
+from .fcn import FCN
 from .lenet import LeNet
 from .metrics import accuracy, accuracy_ignore_pad, cer, wer
 from .mlp import MLP
 from .residual_network import ResidualNetwork, ResidualNetworkEncoder
 from .transformer import Transformer
+from .unet import UNet
 from .util import sliding_window
 from .wide_resnet import WideResidualNetwork
 
@@ -18,12 +20,14 @@ __all__ = [
     "CNNTransformer",
     "ConvolutionalRecurrentNetwork",
     "DenseNet",
+    "FCN",
     "greedy_decoder",
     "MLP",
     "LeNet",
     "ResidualNetwork",
     "ResidualNetworkEncoder",
     "sliding_window",
+    "UNet",
     "Transformer",
     "wer",
     "WideResidualNetwork",
diff --git a/src/text_recognizer/networks/beam.py b/src/text_recognizer/networks/beam.py
new file mode 100644
index 0000000..dccccdb
--- /dev/null
+++ b/src/text_recognizer/networks/beam.py
@@ -0,0 +1,83 @@
+"""Implementation of beam search decoder for a sequence to sequence network.
+
+Stolen from: https://github.com/budzianowski/PyTorch-Beam-Search-Decoding/blob/master/decode_beam.py
+
+"""
+# from typing import List
+# from Queue import PriorityQueue
+
+# from loguru import logger
+# import torch
+# from torch import nn
+# from torch import Tensor
+# import torch.nn.functional as F
+
+
+# class Node:
+#     def __init__(
+#         self, parent: Node, target_index: int, log_prob: Tensor, length: int
+#     ) -> None:
+#         self.parent = parent
+#         self.target_index = target_index
+#         self.log_prob = log_prob
+#         self.length = length
+#         self.reward = 0.0
+
+#     def eval(self, alpha: float = 1.0) -> Tensor:
+#         return self.log_prob / (self.length - 1 + 1e-6) + alpha * self.reward
+
+
+# @torch.no_grad()
+# def beam_decoder(
+#     network, mapper, device, memory: Tensor = None, max_len: int = 97,
+# ) -> Tensor:
+#     beam_width = 10
+#     topk = 1  # How many sentences to generate.
+
+#     trg_indices = [mapper(mapper.init_token)]
+
+#     end_nodes = []
+
+#     node = Node(None, trg_indices, 0, 1)
+#     nodes = PriorityQueue()
+
+#     nodes.put((node.eval(), node))
+#     q_size = 1
+
+#     # Beam search
+#     for _ in range(max_len):
+#         if q_size > 2000:
+#             logger.warning("Could not decoder input")
+#             break
+
+#         # Fetch the best node.
+#         score, n = nodes.get()
+#         decoder_input = n.target_index
+
+#         if n.target_index == mapper(mapper.eos_token) and n.parent is not None:
+#             end_nodes.append((score, n))
+
+#             # If we reached the maximum number of sentences required.
+#             if len(end_nodes) >= 1:
+#                 break
+#             else:
+#                 continue
+
+#         # Forward pass with transformer.
+#         trg = torch.tensor(trg_indices, device=device)[None, :].long()
+#         trg = network.target_embedding(trg)
+#         logits = network.decoder(trg=trg, memory=memory, trg_mask=None)
+#         log_prob = F.log_softmax(logits, dim=2)
+
+#         log_prob, indices = torch.topk(log_prob, beam_width)
+
+#         for new_k in range(beam_width):
+#             # TODO: continue from here
+#             token_index = indices[0][new_k].view(1, -1)
+#             log_p = log_prob[0][new_k].item()
+
+#             node = Node()
+
+#             pass
+
+#     pass
diff --git a/src/text_recognizer/networks/cnn_transformer.py b/src/text_recognizer/networks/cnn_transformer.py
index 16c7a41..b2b74b3 100644
--- a/src/text_recognizer/networks/cnn_transformer.py
+++ b/src/text_recognizer/networks/cnn_transformer.py
@@ -88,10 +88,14 @@ class CNNTransformer(nn.Module):
         if len(src.shape) < 4:
             src = src[(None,) * (4 - len(src.shape))]
         src = self.backbone(src)
-        src = rearrange(src, "b c h w -> b w c h")
+
         if self.adaptive_pool is not None:
+            src = rearrange(src, "b c h w -> b w c h")
             src = self.adaptive_pool(src)
-        src = src.squeeze(3)
+            src = src.squeeze(3)
+        else:
+            src = rearrange(src, "b c h w -> b (w h) c")
+
         src = self.position_encoding(src)
 
         return src
@@ -110,12 +114,17 @@ class CNNTransformer(nn.Module):
         trg = self.position_encoding(trg)
         return trg
 
-    def forward(self, x: Tensor, trg: Optional[Tensor] = None) -> Tensor:
-        """Forward pass with CNN transfomer."""
-        h = self.extract_image_features(x)
+    def decode_image_features(self, h: Tensor, trg: Optional[Tensor] = None) -> Tensor:
+        """Takes images features from the backbone and decodes them with the transformer."""
         trg_mask = self._create_trg_mask(trg)
         trg = self.target_embedding(trg)
         out = self.transformer(h, trg, trg_mask=trg_mask)
 
         logits = self.head(out)
         return logits
+
+    def forward(self, x: Tensor, trg: Optional[Tensor] = None) -> Tensor:
+        """Forward pass with CNN transfomer."""
+        h = self.extract_image_features(x)
+        logits = self.decode_image_features(h, trg)
+        return logits
diff --git a/src/text_recognizer/networks/fcn.py b/src/text_recognizer/networks/fcn.py
deleted file mode 100644
index f9c4fd4..0000000
--- a/src/text_recognizer/networks/fcn.py
+++ /dev/null
@@ -1,99 +0,0 @@
-"""Fully Convolutional Network (FCN) with dilated kernels for global context."""
-from typing import List, Tuple, Type
-import torch
-from torch import nn
-from torch import Tensor
-
-
-from text_recognizer.networks.util import activation_function
-
-
-class _DilatedBlock(nn.Module):
-    def __init__(
-        self,
-        channels: List[int],
-        kernel_sizes: List[int],
-        dilations: List[int],
-        paddings: List[int],
-        activation_fn: Type[nn.Module],
-    ) -> None:
-        super().__init__()
-        self.dilation_conv = nn.Sequential(
-            nn.Conv2d(
-                in_channels=channels[0],
-                out_channels=channels[1],
-                kernel_size=kernel_sizes[0],
-                stride=1,
-                dilation=dilations[0],
-                padding=paddings[0],
-            ),
-            nn.Conv2d(
-                in_channels=channels[1],
-                out_channels=channels[1] // 2,
-                kernel_size=kernel_sizes[1],
-                stride=1,
-                dilation=dilations[1],
-                padding=paddings[1],
-            ),
-        )
-        self.activation_fn = activation_fn
-
-        self.conv = nn.Conv2d(
-            in_channels=channels[0],
-            out_channels=channels[1] // 2,
-            kernel_size=1,
-            dilation=1,
-            stride=1,
-        )
-
-    def forward(self, x: Tensor) -> Tensor:
-        residual = self.conv(x)
-        x = self.dilation_conv(x)
-        x = torch.cat((x, residual), dim=1)
-        return self.activation_fn(x)
-
-
-class FCN(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        base_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        dilations: Tuple[int] = (3, 7),
-        paddings: Tuple[int] = (9, 21),
-        num_blocks: int = 14,
-        activation: str = "elu",
-    ) -> None:
-        super().__init__()
-        self.kernel_sizes = [kernel_size] * num_blocks
-        self.channels = [in_channels] + [base_channels] * (num_blocks - 1)
-        self.out_channels = out_channels
-        self.dilations = [dilations[0]] * (num_blocks // 2) + [dilations[1]] * (
-            num_blocks // 2
-        )
-        self.paddings = [paddings[0]] * (num_blocks // 2) + [paddings[1]] * (
-            num_blocks // 2
-        )
-        self.activation_fn = activation_function(activation)
-        self.fcn = self._configure_fcn()
-
-    def _configure_fcn(self) -> nn.Sequential:
-        layers = []
-        for i in range(0, len(self.channels), 2):
-            layers.append(
-                _DilatedBlock(
-                    self.channels[i : i + 2],
-                    self.kernel_sizes[i : i + 2],
-                    self.dilations[i : i + 2],
-                    self.paddings[i : i + 2],
-                    self.activation_fn,
-                )
-            )
-        layers.append(
-            nn.Conv2d(self.channels[-1], self.out_channels, kernel_size=1, stride=1)
-        )
-        return nn.Sequential(*layers)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.fcn(x)
diff --git a/src/text_recognizer/networks/neural_machine_reader.py b/src/text_recognizer/networks/neural_machine_reader.py
deleted file mode 100644
index 7f8c49b..0000000
--- a/src/text_recognizer/networks/neural_machine_reader.py
+++ /dev/null
@@ -1,201 +0,0 @@
-"""Sequence to sequence network with RNN cells."""
-# from typing import Dict, Optional, Tuple
-
-# from einops import rearrange
-# from einops.layers.torch import Rearrange
-# import torch
-# from torch import nn
-# from torch import Tensor
-
-# from text_recognizer.networks.util import configure_backbone
-
-
-# class Encoder(nn.Module):
-#     def __init__(
-#         self,
-#         embedding_dim: int,
-#         encoder_dim: int,
-#         decoder_dim: int,
-#         dropout_rate: float = 0.1,
-#     ) -> None:
-#         super().__init__()
-#         self.rnn = nn.GRU(
-#             input_size=embedding_dim, hidden_size=encoder_dim, bidirectional=True
-#         )
-#         self.fc = nn.Sequential(
-#             nn.Linear(in_features=2 * encoder_dim, out_features=decoder_dim), nn.Tanh()
-#         )
-#         self.dropout = nn.Dropout(p=dropout_rate)
-
-#     def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
-#         """Encodes a sequence of tensors with a bidirectional GRU.
-
-#         Args:
-#             x (Tensor): A input sequence.
-
-#         Shape:
-#             - x: :math:`(T, N, E)`.
-#             - output[0]: :math:`(T, N, 2 * E)`.
-#             - output[1]: :math:`(T, N, D)`.
-
-#             where T is the sequence length, N is the batch size, E is the
-#             embedding/encoder dimension, and D is the decoder dimension.
-
-#         Returns:
-#             Tuple[Tensor, Tensor]: The encoder output and the hidden state of the
-#                 encoder.
-
-#         """
-
-#         output, hidden = self.rnn(x)
-
-#         # Get the hidden state from the forward and backward rnn.
-#         hidden_state = torch.cat((hidden[-2, :, :], hidden[-1, :, :]), dim=1)
-
-#         # Apply fully connected layer and tanh activation.
-#         hidden_state = self.fc(hidden_state)
-
-#         return output, hidden_state
-
-
-# class Attention(nn.Module):
-#     def __init__(self, encoder_dim: int, decoder_dim: int) -> None:
-#         super().__init__()
-#         self.atten = nn.Linear(
-#             in_features=2 * encoder_dim + decoder_dim, out_features=decoder_dim
-#         )
-#         self.value = nn.Linear(in_features=decoder_dim, out_features=1, bias=False)
-
-#     def forward(self, hidden_state: Tensor, encoder_outputs: Tensor) -> Tensor:
-#         """Short summary.
-
-#         Args:
-#             hidden_state (Tensor): Description of parameter `h`.
-#             encoder_outputs (Tensor): Description of parameter `enc_out`.
-
-#         Shape:
-#             - x: :math:`(T, N, E)`.
-#             - output[0]: :math:`(T, N, 2 * E)`.
-#             - output[1]: :math:`(T, N, D)`.
-
-#             where T is the sequence length, N is the batch size, E is the
-#             embedding/encoder dimension, and D is the decoder dimension.
-
-#         Returns:
-#             Tensor: Description of returned object.
-
-#         """
-#         t, b = enc_out.shape[:2]
-#         # repeat decoder hidden state src_len times
-#         hidden_state = hidden_state.unsqueeze(1).repeat(1, t, 1)
-
-#         encoder_outputs = rearrange(encoder_outputs, "t b e2 -> b t e2")
-
-#         # Calculate the energy between the decoders previous hidden state and the
-#         # encoders hidden states.
-#         energy = torch.tanh(
-#             self.attn(torch.cat((hidden_state, encoder_outputs), dim=2))
-#         )
-
-#         attention = self.value(energy).squeeze(2)
-
-#         # Apply softmax on the attention to squeeze it between 0 and 1.
-#         attention = F.softmax(attention, dim=1)
-
-#         return attention
-
-
-# class Decoder(nn.Module):
-#     def __init__(
-#         self,
-#         embedding_dim: int,
-#         encoder_dim: int,
-#         decoder_dim: int,
-#         output_dim: int,
-#         dropout_rate: float = 0.1,
-#     ) -> None:
-#         super().__init__()
-#         self.output_dim = output_dim
-#         self.embedding = nn.Embedding(output_dim, embedding_dim)
-#         self.attention = Attention(encoder_dim, decoder_dim)
-#         self.rnn = nn.GRU(
-#             input_size=2 * encoder_dim + embedding_dim, hidden_size=decoder_dim
-#         )
-
-#         self.head = nn.Linear(
-#             in_features=2 * encoder_dim + embedding_dim + decoder_dim,
-#             out_features=output_dim,
-#         )
-#         self.dropout = nn.Dropout(p=dropout_rate)
-
-#     def forward(
-#         self, trg: Tensor, hidden_state: Tensor, encoder_outputs: Tensor
-#     ) -> Tensor:
-#         # input = [batch size]
-#         # hidden = [batch size, dec hid dim]
-#         # encoder_outputs = [src len, batch size, enc hid dim * 2]
-#         trg = trg.unsqueeze(0)
-#         trg_embedded = self.dropout(self.embedding(trg))
-
-#         a = self.attention(hidden_state, encoder_outputs)
-
-#         weighted = torch.bmm(a, encoder_outputs)
-
-#         # Permutate the tensor.
-#         weighted = rearrange(weighted, "b a e2 -> a b e2")
-
-#         rnn_input = torch.cat((trg_embedded, weighted), dim=2)
-
-#         output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
-
-#         # seq len, n layers and n directions will always be 1 in this decoder, therefore:
-#         # output = [1, batch size, dec hid dim]
-#         # hidden = [1, batch size, dec hid dim]
-#         # this also means that output == hidden
-#         assert (output == hidden).all()
-
-#         trg_embedded = trg_embedded.squeeze(0)
-#         output = output.squeeze(0)
-#         weighted = weighted.squeeze(0)
-
-#         logits = self.fc_out(torch.cat((output, weighted, trg_embedded), dim=1))
-
-#         # prediction = [batch size, output dim]
-
-#         return logits, hidden.squeeze(0)
-
-
-# class NeuralMachineReader(nn.Module):
-#     def __init__(
-#         self,
-#         embedding_dim: int,
-#         encoder_dim: int,
-#         decoder_dim: int,
-#         output_dim: int,
-#         backbone: Optional[str] = None,
-#         backbone_args: Optional[Dict] = None,
-#         adaptive_pool_dim: Tuple = (None, 1),
-#         dropout_rate: float = 0.1,
-#         teacher_forcing_ratio: float = 0.5,
-#     ) -> None:
-#         super().__init__()
-
-#         self.backbone = configure_backbone(backbone, backbone_args)
-#         self.adaptive_pool = nn.AdaptiveAvgPool2d((adaptive_pool_dim))
-
-#         self.encoder = Encoder(embedding_dim, encoder_dim, decoder_dim, dropout_rate)
-#         self.decoder = Decoder(
-#             embedding_dim, encoder_dim, decoder_dim, output_dim, dropout_rate
-#         )
-#         self.teacher_forcing_ratio = teacher_forcing_ratio
-
-#     def extract_image_features(self, x: Tensor) -> Tensor:
-#         x = self.backbone(x)
-#         x = rearrange(x, "b c h w -> b w c h")
-#         x = self.adaptive_pool(x)
-#         x = x.squeeze(3)
-
-#     def forward(self, x: Tensor, trg: Tensor) -> Tensor:
-#         # x = [batch size, height, width]
-#         # trg = [trg len, batch size]
-#         z = self.extract_image_features(x)
diff --git a/src/text_recognizer/networks/residual_network.py b/src/text_recognizer/networks/residual_network.py
index 6405192..e397224 100644
--- a/src/text_recognizer/networks/residual_network.py
+++ b/src/text_recognizer/networks/residual_network.py
@@ -7,7 +7,6 @@ import torch
 from torch import nn
 from torch import Tensor
 
-from text_recognizer.networks.stn import SpatialTransformerNetwork
 from text_recognizer.networks.util import activation_function
 
 
@@ -209,12 +208,10 @@ class ResidualNetworkEncoder(nn.Module):
         activation: str = "relu",
         block: Type[nn.Module] = BasicBlock,
         levels: int = 1,
-        stn: bool = False,
         *args,
         **kwargs
     ) -> None:
         super().__init__()
-        self.stn = SpatialTransformerNetwork() if stn else None
         self.block_sizes = (
             block_sizes if isinstance(block_sizes, list) else [block_sizes] * levels
         )
@@ -231,7 +228,7 @@ class ResidualNetworkEncoder(nn.Module):
             ),
             nn.BatchNorm2d(self.block_sizes[0]),
             activation_function(self.activation),
-            nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
+            # nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
         )
 
         self.blocks = self._configure_blocks(block)
@@ -275,8 +272,6 @@ class ResidualNetworkEncoder(nn.Module):
         # If batch dimenstion is missing, it needs to be added.
         if len(x.shape) == 3:
             x = x.unsqueeze(0)
-        if self.stn is not None:
-            x = self.stn(x)
         x = self.gate(x)
         x = self.blocks(x)
         return x
diff --git a/src/text_recognizer/networks/unet.py b/src/text_recognizer/networks/unet.py
index 51f242a..510910f 100644
--- a/src/text_recognizer/networks/unet.py
+++ b/src/text_recognizer/networks/unet.py
@@ -8,64 +8,118 @@ from torch import Tensor
 from text_recognizer.networks.util import activation_function
 
 
-class ConvBlock(nn.Module):
-    """Basic UNet convolutional block."""
+class _ConvBlock(nn.Module):
+    """Modified UNet convolutional block with dilation."""
 
-    def __init__(self, channels: List[int], activation: str) -> None:
+    def __init__(
+        self,
+        channels: List[int],
+        activation: str,
+        num_groups: int,
+        dropout_rate: float = 0.1,
+        kernel_size: int = 3,
+        dilation: int = 1,
+        padding: int = 0,
+    ) -> None:
         super().__init__()
         self.channels = channels
+        self.dropout_rate = dropout_rate
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.padding = padding
+        self.num_groups = num_groups
         self.activation = activation_function(activation)
         self.block = self._configure_block()
+        self.residual_conv = nn.Sequential(
+            nn.Conv2d(
+                self.channels[0], self.channels[-1], kernel_size=3, stride=1, padding=1
+            ),
+            self.activation,
+        )
 
     def _configure_block(self) -> nn.Sequential:
         block = []
         for i in range(len(self.channels) - 1):
             block += [
+                nn.Dropout(p=self.dropout_rate),
+                nn.GroupNorm(self.num_groups, self.channels[i]),
+                self.activation,
                 nn.Conv2d(
-                    self.channels[i], self.channels[i + 1], kernel_size=3, padding=1
+                    self.channels[i],
+                    self.channels[i + 1],
+                    kernel_size=self.kernel_size,
+                    padding=self.padding,
+                    stride=1,
+                    dilation=self.dilation,
                 ),
-                nn.BatchNorm2d(self.channels[i + 1]),
-                self.activation,
             ]
 
         return nn.Sequential(*block)
 
     def forward(self, x: Tensor) -> Tensor:
         """Apply the convolutional block."""
-        return self.block(x)
+        residual = self.residual_conv(x)
+        return self.block(x) + residual
 
 
-class DownSamplingBlock(nn.Module):
+class _DownSamplingBlock(nn.Module):
     """Basic down sampling block."""
 
     def __init__(
         self,
         channels: List[int],
         activation: str,
+        num_groups: int,
         pooling_kernel: Union[int, bool] = 2,
+        dropout_rate: float = 0.1,
+        kernel_size: int = 3,
+        dilation: int = 1,
+        padding: int = 0,
     ) -> None:
         super().__init__()
-        self.conv_block = ConvBlock(channels, activation)
+        self.conv_block = _ConvBlock(
+            channels,
+            activation,
+            num_groups,
+            dropout_rate,
+            kernel_size,
+            dilation,
+            padding,
+        )
         self.down_sampling = nn.MaxPool2d(pooling_kernel) if pooling_kernel else None
 
     def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
         """Return the convolutional block output and a down sampled tensor."""
         x = self.conv_block(x)
-        if self.down_sampling is not None:
-            x_down = self.down_sampling(x)
-        else:
-            x_down = None
+        x_down = self.down_sampling(x) if self.down_sampling is not None else x
+
         return x_down, x
 
 
-class UpSamplingBlock(nn.Module):
+class _UpSamplingBlock(nn.Module):
     """The upsampling block of the UNet."""
 
     def __init__(
-        self, channels: List[int], activation: str, scale_factor: int = 2
+        self,
+        channels: List[int],
+        activation: str,
+        num_groups: int,
+        scale_factor: int = 2,
+        dropout_rate: float = 0.1,
+        kernel_size: int = 3,
+        dilation: int = 1,
+        padding: int = 0,
     ) -> None:
         super().__init__()
-        self.conv_block = ConvBlock(channels, activation)
+        self.conv_block = _ConvBlock(
+            channels,
+            activation,
+            num_groups,
+            dropout_rate,
+            kernel_size,
+            dilation,
+            padding,
+        )
         self.up_sampling = nn.Upsample(
             scale_factor=scale_factor, mode="bilinear", align_corners=True
         )
@@ -87,14 +141,43 @@ class UNet(nn.Module):
         base_channels: int = 64,
         num_classes: int = 3,
         depth: int = 4,
-        out_channels: int = 3,
         activation: str = "relu",
+        num_groups: int = 8,
+        dropout_rate: float = 0.1,
         pooling_kernel: int = 2,
         scale_factor: int = 2,
+        kernel_size: Optional[List[int]] = None,
+        dilation: Optional[List[int]] = None,
+        padding: Optional[List[int]] = None,
     ) -> None:
         super().__init__()
         self.depth = depth
-        channels = [1] + [base_channels * 2 ** i for i in range(depth)]
+        self.num_groups = num_groups
+
+        if kernel_size is not None and dilation is not None and padding is not None:
+            if (
+                len(kernel_size) != depth
+                and len(dilation) != depth
+                and len(padding) != depth
+            ):
+                raise RuntimeError(
+                    "Length of convolutional parameters does not match the depth."
+                )
+            self.kernel_size = kernel_size
+            self.padding = padding
+            self.dilation = dilation
+
+        else:
+            self.kernel_size = [3] * depth
+            self.padding = [1] * depth
+            self.dilation = [1] * depth
+
+        self.dropout_rate = dropout_rate
+        self.conv = nn.Conv2d(
+            in_channels, base_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        channels = [base_channels] + [base_channels * 2 ** i for i in range(depth)]
         self.encoder_blocks = self._configure_down_sampling_blocks(
             channels, activation, pooling_kernel
         )
@@ -110,49 +193,63 @@ class UNet(nn.Module):
         blocks = nn.ModuleList([])
         for i in range(len(channels) - 1):
             pooling_kernel = pooling_kernel if i < self.depth - 1 else False
+            dropout_rate = self.dropout_rate if i < 0 else 0
             blocks += [
-                DownSamplingBlock(
+                _DownSamplingBlock(
                     [channels[i], channels[i + 1], channels[i + 1]],
                     activation,
+                    self.num_groups,
                     pooling_kernel,
+                    dropout_rate,
+                    self.kernel_size[i],
+                    self.dilation[i],
+                    self.padding[i],
                 )
             ]
 
         return blocks
 
     def _configure_up_sampling_blocks(
-        self,
-        channels: List[int],
-        activation: str,
-        scale_factor: int,
+        self, channels: List[int], activation: str, scale_factor: int,
     ) -> nn.ModuleList:
         channels.reverse()
+        self.kernel_size.reverse()
+        self.dilation.reverse()
+        self.padding.reverse()
         return nn.ModuleList(
             [
-                UpSamplingBlock(
+                _UpSamplingBlock(
                     [channels[i] + channels[i + 1], channels[i + 1], channels[i + 1]],
                     activation,
+                    self.num_groups,
                     scale_factor,
+                    self.dropout_rate,
+                    self.kernel_size[i],
+                    self.dilation[i],
+                    self.padding[i],
                 )
                 for i in range(len(channels) - 2)
             ]
         )
 
-    def encode(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
+    def _encode(self, x: Tensor) -> List[Tensor]:
         x_skips = []
         for block in self.encoder_blocks:
             x, x_skip = block(x)
-            if x_skip is not None:
-                x_skips.append(x_skip)
-        return x, x_skips
+            x_skips.append(x_skip)
+        return x_skips
 
-    def decode(self, x: Tensor, x_skips: List[Tensor]) -> Tensor:
+    def _decode(self, x_skips: List[Tensor]) -> Tensor:
         x = x_skips[-1]
         for i, block in enumerate(self.decoder_blocks):
             x = block(x, x_skips[-(i + 2)])
         return x
 
     def forward(self, x: Tensor) -> Tensor:
-        x, x_skips = self.encode(x)
-        x = self.decode(x, x_skips)
+        """Forward pass with the UNet model."""
+        if len(x.shape) < 4:
+            x = x[(None,) * (4 - len(x.shape))]
+        x = self.conv(x)
+        x_skips = self._encode(x)
+        x = self._decode(x_skips)
         return self.head(x)
diff --git a/src/text_recognizer/paragraph_text_recognizer.py b/src/text_recognizer/paragraph_text_recognizer.py
new file mode 100644
index 0000000..aa39662
--- /dev/null
+++ b/src/text_recognizer/paragraph_text_recognizer.py
@@ -0,0 +1,153 @@
+"""Full model.
+
+Takes an image and returns the text in the image, by first segmenting the image with a LineDetector, then extracting the
+each crop of the image corresponding to line regions, and feeding them to a LinePredictor model that outputs the text
+in each region.
+"""
+from typing import Dict, List, Tuple, Union
+
+import cv2
+import numpy as np
+import torch
+
+from text_recognizer.models import SegmentationModel, TransformerModel
+from text_recognizer.util import read_image
+
+
+class ParagraphTextRecognizor:
+    """Given an image of a single handwritten character, recognizes it."""
+
+    def __init__(self, line_predictor_args: Dict, line_detector_args: Dict) -> None:
+        self._line_predictor = TransformerModel(**line_predictor_args)
+        self._line_detector = SegmentationModel(**line_detector_args)
+        self._line_detector.eval()
+        self._line_predictor.eval()
+
+    def predict(self, image_or_filename: Union[str, np.ndarray]) -> Tuple:
+        """Takes an image and returns all text within it."""
+        image = (
+            read_image(image_or_filename)
+            if isinstance(image_or_filename, str)
+            else image_or_filename
+        )
+
+        line_region_crops = self._get_line_region_crops(image)
+        processed_line_region_crops = [
+            self._process_image_for_line_predictor(image=crop)
+            for crop in line_region_crops
+        ]
+        line_region_strings = [
+            self.line_predictor_model.predict_on_image(crop)[0]
+            for crop in processed_line_region_crops
+        ]
+
+        return " ".join(line_region_strings), line_region_crops
+
+    def _get_line_region_crops(
+        self, image: np.ndarray, min_crop_len_factor: float = 0.02
+    ) -> List[np.ndarray]:
+        """Returns all the crops of text lines in a square image."""
+        processed_image, scale_down_factor = self._process_image_for_line_detector(
+            image
+        )
+        line_segmentation = self._line_detector.predict_on_image(processed_image)
+        bounding_boxes = _find_line_bounding_boxes(line_segmentation)
+
+        bounding_boxes = (bounding_boxes * scale_down_factor).astype(int)
+
+        min_crop_len = int(min_crop_len_factor * min(image.shape[0], image.shape[1]))
+        line_region_crops = [
+            image[y : y + h, x : x + w]
+            for x, y, w, h in bounding_boxes
+            if w >= min_crop_len and h >= min_crop_len
+        ]
+        return line_region_crops
+
+    def _process_image_for_line_detector(
+        self, image: np.ndarray
+    ) -> Tuple[np.ndarray, float]:
+        """Convert uint8 image to float image with black background with shape self._line_detector.image_shape."""
+        resized_image, scale_down_factor = _resize_image_for_line_detector(
+            image=image, max_shape=self._line_detector.image_shape
+        )
+        resized_image = (1.0 - resized_image / 255).astype("float32")
+        return resized_image, scale_down_factor
+
+    def _process_image_for_line_predictor(self, image: np.ndarray) -> np.ndarray:
+        """Preprocessing of image before feeding it to the LinePrediction model.
+
+        Convert uint8 image to float image with black background with shape
+        self._line_predictor.image_shape while maintaining the image aspect ratio.
+
+        Args:
+            image (np.ndarray): Crop of text line.
+
+        Returns:
+            np.ndarray: Processed crop for feeding line predictor.
+        """
+        expected_shape = self._line_detector.image_shape
+        scale_factor = (np.array(expected_shape) / np.array(image.shape)).min()
+        scaled_image = cv2.resize(
+            image,
+            dsize=None,
+            fx=scale_factor,
+            fy=scale_factor,
+            interpolation=cv2.INTER_AREA,
+        )
+
+        pad_with = (
+            (0, expected_shape[0] - scaled_image.shape[0]),
+            (0, expected_shape[1] - scaled_image.shape[1]),
+        )
+
+        padded_image = np.pad(
+            scaled_image, pad_with=pad_with, mode="constant", constant_values=255
+        )
+        return 1 - padded_image / 255
+
+
+def _find_line_bounding_boxes(line_segmentation: np.ndarray) -> np.ndarray:
+    """Given a line segmentation, find bounding boxes for connected-component regions corresponding to non-0 labels."""
+
+    def _find_line_bounding_boxes_in_channel(
+        line_segmentation_channel: np.ndarray,
+    ) -> np.ndarray:
+        line_segmentation_image = cv2.dilate(
+            line_segmentation_channel, kernel=np.ones((3, 3)), iterations=1
+        )
+        line_activation_image = (line_segmentation_image * 255).astype("uint8")
+        line_activation_image = cv2.threshold(
+            line_activation_image, 0.5, 1, cv2.THRESH_BINARY | cv2.THRESH_OTSU
+        )[1]
+
+        bounding_cnts, _ = cv2.findContours(
+            line_segmentation_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+        return np.array([cv2.boundingRect(cnt) for cnt in bounding_cnts])
+
+    bounding_boxes = np.concatenate(
+        [
+            _find_line_bounding_boxes_in_channel(line_segmentation[:, :, i])
+            for i in [1, 2]
+        ],
+        axis=0,
+    )
+
+    return bounding_boxes[np.argsort(bounding_boxes[:, 1])]
+
+
+def _resize_image_for_line_detector(
+    image: np.ndarray, max_shape: Tuple[int, int]
+) -> Tuple[np.ndarray, float]:
+    """Resize the image to less than the max_shape while maintaining the aspect ratio."""
+    scale_down_factor = max(np.ndarray(image.shape) / np.ndarray(max_shape))
+    if scale_down_factor == 1:
+        return image.copy(), scale_down_factor
+    resize_image = cv2.resize(
+        image,
+        dsize=None,
+        fx=1 / scale_down_factor,
+        fy=1 / scale_down_factor,
+        interpolation=cv2.INTER_AREA,
+    )
+    return resize_image, scale_down_factor
diff --git a/src/text_recognizer/tests/support/iam_paragraphs/a01-000u.jpg b/src/text_recognizer/tests/support/iam_paragraphs/a01-000u.jpg
new file mode 100644
index 0000000..d9753b6
--- /dev/null
+++ b/src/text_recognizer/tests/support/iam_paragraphs/a01-000u.jpg
diff --git a/src/text_recognizer/tests/test_paragraph_text_recognizer.py b/src/text_recognizer/tests/test_paragraph_text_recognizer.py
new file mode 100644
index 0000000..3e280b9
--- /dev/null
+++ b/src/text_recognizer/tests/test_paragraph_text_recognizer.py
@@ -0,0 +1,37 @@
+"""Test for ParagraphTextRecognizer class."""
+import os
+from pathlib import Path
+import unittest
+
+from text_recognizer.paragraph_text_recognizer import ParagraphTextRecognizor
+import text_recognizer.util as util
+
+
+SUPPORT_DIRNAME = Path(__file__).parents[0].resolve() / "support" / "iam_paragraph"
+
+# Prevent using GPU.
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+
+class TestParagraphTextRecognizor(unittest.TestCase):
+    """Test that it can take non-square images of max dimension larger than 256px."""
+
+    def test_filename(self) -> None:
+        """Test model on support image."""
+        line_predictor_args = {
+            "dataset": "EmnistLineDataset",
+            "network_fn": "CNNTransformer",
+        }
+        line_detector_args = {"dataset": "EmnistLineDataset", "network_fn": "UNet"}
+        model = ParagraphTextRecognizor(
+            line_predictor_args=line_predictor_args,
+            line_detector_args=line_detector_args,
+        )
+        num_text_lines_by_name = {"a01-000u-cropped": 7}
+        for filename in (SUPPORT_DIRNAME).glob("*.jpg"):
+            full_image = util.read_image(str(filename), grayscale=True)
+            predicted_text, line_region_crops = model.predict(full_image)
+            print(predicted_text)
+            self.assertTrue(
+                len(line_region_crops), num_text_lines_by_name[filename.stem]
+            )
diff --git a/src/text_recognizer/util.py b/src/text_recognizer/util.py
index 6c07c60..b431e22 100644
--- a/src/text_recognizer/util.py
+++ b/src/text_recognizer/util.py
@@ -21,20 +21,21 @@ def read_image(image_uri: Union[Path, str], grayscale: bool = False) -> np.ndarr
             return cv2.imdecode(image_array, imread_flag)
         else:
             raise ValueError(
-                "Url does not start with http, therfore not safe to open..."
+                "Url does not start with http, therefore not safe to open..."
             ) from None
 
     imread_flag = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR
     local_file = os.path.exists(image_uri)
-    try:
-        image = None
-        if local_file:
-            image = read_image_from_filename(image_uri, imread_flag)
-        else:
-            image = read_image_from_url(image_uri, imread_flag)
-        assert image is not None
-    except Exception as e:
-        raise ValueError(f"Could not load image at {image_uri}: {e}")
+    image = None
+
+    if local_file:
+        image = read_image_from_filename(image_uri, imread_flag)
+    else:
+        image = read_image_from_url(image_uri, imread_flag)
+
+    if image is None:
+        raise ValueError(f"Could not load image at {image_uri}")
+
     return image
 
 
diff --git a/src/text_recognizer/weights/SegmentationModel_IamParagraphsDataset_FCN_weights.pt b/src/text_recognizer/weights/SegmentationModel_IamParagraphsDataset_FCN_weights.pt
new file mode 100644
index 0000000..d9ca01d
--- /dev/null
+++ b/src/text_recognizer/weights/SegmentationModel_IamParagraphsDataset_FCN_weights.pt
diff --git a/src/text_recognizer/weights/SegmentationModel_IamParagraphsDataset_UNet_weights.pt b/src/text_recognizer/weights/SegmentationModel_IamParagraphsDataset_UNet_weights.pt
new file mode 100644
index 0000000..0af0e57
--- /dev/null
+++ b/src/text_recognizer/weights/SegmentationModel_IamParagraphsDataset_UNet_weights.pt