diff options
Diffstat (limited to 'src/text_recognizer')
18 files changed, 263 insertions, 580 deletions
diff --git a/src/text_recognizer/datasets/transforms.py b/src/text_recognizer/datasets/transforms.py index 8deac7f..1105f23 100644 --- a/src/text_recognizer/datasets/transforms.py +++ b/src/text_recognizer/datasets/transforms.py @@ -3,7 +3,8 @@ import numpy as np from PIL import Image import torch from torch import Tensor -from torchvision.transforms import Compose, Resize, ToPILImage, ToTensor +import torch.nn.functional as F +from torchvision.transforms import Compose, ToPILImage, ToTensor from text_recognizer.datasets.util import EmnistMapper @@ -16,6 +17,18 @@ class Transpose: return np.array(image).swapaxes(0, 1) +class Resize: + """Resizes a tensor to a specified width.""" + + def __init__(self, width: int = 952) -> None: + # The default is 952 because of the IAM dataset. + self.width = width + + def __call__(self, image: Tensor) -> Tensor: + """Resize tensor in the last dimension.""" + return F.interpolate(image, size=self.width, mode="nearest") + + class AddTokens: """Adds start of sequence and end of sequence tokens to target tensor.""" diff --git a/src/text_recognizer/models/__init__.py b/src/text_recognizer/models/__init__.py index 28aa52e..53340f1 100644 --- a/src/text_recognizer/models/__init__.py +++ b/src/text_recognizer/models/__init__.py @@ -2,18 +2,16 @@ from .base import Model from .character_model import CharacterModel from .crnn_model import CRNNModel -from .metrics import accuracy, cer, wer -from .transformer_encoder_model import TransformerEncoderModel -from .vision_transformer_model import VisionTransformerModel +from .metrics import accuracy, accuracy_ignore_pad, cer, wer +from .transformer_model import TransformerModel __all__ = [ - "Model", + "accuracy", + "accuracy_ignore_pad", "cer", "CharacterModel", "CRNNModel", - "CNNTransfromerModel", - "accuracy", - "TransformerEncoderModel", - "VisionTransformerModel", + "Model", + "TransformerModel", "wer", ] diff --git a/src/text_recognizer/models/base.py b/src/text_recognizer/models/base.py index cc44c92..a945b41 100644 --- a/src/text_recognizer/models/base.py +++ b/src/text_recognizer/models/base.py @@ -49,7 +49,7 @@ class Model(ABC): network_args (Optional[Dict]): Arguments for the network. Defaults to None. dataset_args (Optional[Dict]): Arguments for the dataset. metrics (Optional[Dict]): Metrics to evaluate the performance with. Defaults to None. - criterion (Optional[Callable]): The criterion to evaulate the preformance of the network. + criterion (Optional[Callable]): The criterion to evaluate the performance of the network. Defaults to None. criterion_args (Optional[Dict]): Dict of arguments for criterion. Defaults to None. optimizer (Optional[Callable]): The optimizer for updating the weights. Defaults to None. @@ -221,7 +221,7 @@ class Model(ABC): def _configure_network(self, network_fn: Type[nn.Module]) -> None: """Loads the network.""" - # If no network arguemnts are given, load pretrained weights if they exist. + # If no network arguments are given, load pretrained weights if they exist. if self._network_args is None: self.load_weights(network_fn) else: @@ -245,7 +245,7 @@ class Model(ABC): self._optimizer = None if self._optimizer and self._lr_scheduler is not None: - if "OneCycleLR" in str(self._lr_scheduler): + if "steps_per_epoch" in self.lr_scheduler_args: self.lr_scheduler_args["steps_per_epoch"] = len(self.train_dataloader()) # Assume lr scheduler should update at each epoch if not specified. @@ -412,7 +412,7 @@ class Model(ABC): self._optimizer.load_state_dict(checkpoint["optimizer_state"]) if self._lr_scheduler is not None: - # Does not work when loadning from previous checkpoint and trying to train beyond the last max epochs + # Does not work when loading from previous checkpoint and trying to train beyond the last max epochs # with OneCycleLR. if self._lr_scheduler["lr_scheduler"].__class__.__name__ != "OneCycleLR": self._lr_scheduler["lr_scheduler"].load_state_dict( diff --git a/src/text_recognizer/models/metrics.py b/src/text_recognizer/models/metrics.py index 42c3c6e..af9adb5 100644 --- a/src/text_recognizer/models/metrics.py +++ b/src/text_recognizer/models/metrics.py @@ -6,7 +6,23 @@ from torch import Tensor from text_recognizer.networks import greedy_decoder -def accuracy(outputs: Tensor, labels: Tensor) -> float: +def accuracy_ignore_pad( + output: Tensor, + target: Tensor, + pad_index: int = 79, + eos_index: int = 81, + seq_len: int = 97, +) -> float: + """Sets all predictions after eos to pad.""" + start_indices = torch.nonzero(target == eos_index, as_tuple=False).squeeze(1) + end_indices = torch.arange(seq_len, target.shape[0] + 1, seq_len) + for start, stop in zip(start_indices, end_indices): + output[start + 1 : stop] = pad_index + + return accuracy(output, target) + + +def accuracy(outputs: Tensor, labels: Tensor,) -> float: """Computes the accuracy. Args: @@ -17,10 +33,9 @@ def accuracy(outputs: Tensor, labels: Tensor) -> float: float: The accuracy for the batch. """ - # eos_index = torch.nonzero(labels == eos, as_tuple=False) - # eos_index = eos_index[0].item() if eos_index.nelement() else -1 _, predicted = torch.max(outputs, dim=-1) + acc = (predicted == labels).sum().float() / labels.shape[0] acc = acc.item() return acc diff --git a/src/text_recognizer/models/transformer_encoder_model.py b/src/text_recognizer/models/transformer_encoder_model.py deleted file mode 100644 index e35e298..0000000 --- a/src/text_recognizer/models/transformer_encoder_model.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Defines the CNN-Transformer class.""" -from typing import Callable, Dict, List, Optional, Tuple, 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.datasets import EmnistMapper -from text_recognizer.models.base import Model - - -class TransformerEncoderModel(Model): - """A class for only using the encoder part in the sequence modelling.""" - - def __init__( - self, - network_fn: Type[nn.Module], - dataset: Type[Dataset], - 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.init_token = dataset_args["args"]["init_token"] - self.pad_token = dataset_args["args"]["pad_token"] - self.eos_token = dataset_args["args"]["eos_token"] - if network_args is not None: - self.max_len = network_args["max_len"] - else: - self.max_len = 128 - - if self._mapper is None: - self._mapper = EmnistMapper( - # init_token=self.init_token, - pad_token=self.pad_token, - eos_token=self.eos_token, - ) - self.tensor_transform = ToTensor() - - self.softmax = nn.Softmax(dim=2) - - @torch.no_grad() - def _generate_sentence(self, image: Tensor) -> Tuple[List, float]: - logits = self.network(image) - # Convert logits to probabilities. - probs = self.softmax(logits).squeeze(0) - - confidence, pred_tokens = probs.max(1) - pred_tokens = pred_tokens - - eos_index = torch.nonzero( - pred_tokens == self._mapper(self.eos_token), as_tuple=False, - ) - - eos_index = eos_index[0].item() if eos_index.nelement() else -1 - - predicted_characters = "".join( - [self.mapper(x) for x in pred_tokens[:eos_index].tolist()] - ) - - confidence = np.min(confidence.tolist()) - - return predicted_characters, confidence - - @torch.no_grad() - def predict_on_image(self, image: Union[np.ndarray, Tensor]) -> Tuple[str, float]: - """Predict on a single input.""" - self.eval() - - if image.dtype == 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 == torch.uint8: - # If the image is an unscaled tensor. - image = image.type("torch.FloatTensor") / 255 - - # Put the image tensor on the device the model weights are on. - image = image.to(self.device) - - (predicted_characters, confidence_of_prediction,) = self._generate_sentence( - image - ) - - return predicted_characters, confidence_of_prediction diff --git a/src/text_recognizer/models/vision_transformer_model.py b/src/text_recognizer/models/transformer_model.py index 3d36437..968a047 100644 --- a/src/text_recognizer/models/vision_transformer_model.py +++ b/src/text_recognizer/models/transformer_model.py @@ -13,7 +13,7 @@ from text_recognizer.models.base import Model from text_recognizer.networks import greedy_decoder -class VisionTransformerModel(Model): +class TransformerModel(Model): """Model for predicting a sequence of characters from an image of a text line with a cnn-transformer.""" def __init__( @@ -50,10 +50,7 @@ class VisionTransformerModel(Model): self.init_token = dataset_args["args"]["init_token"] self.pad_token = dataset_args["args"]["pad_token"] self.eos_token = dataset_args["args"]["eos_token"] - if network_args is not None: - self.max_len = network_args["max_len"] - else: - self.max_len = 120 + self.max_len = 120 if self._mapper is None: self._mapper = EmnistMapper( @@ -67,7 +64,7 @@ class VisionTransformerModel(Model): @torch.no_grad() def _generate_sentence(self, image: Tensor) -> Tuple[List, float]: - src = self.network.preprocess_input(image) + src = self.network.extract_image_features(image) memory = self.network.encoder(src) confidence_of_predictions = [] @@ -75,7 +72,7 @@ class VisionTransformerModel(Model): for _ in range(self.max_len - 1): trg = torch.tensor(trg_indices, device=self.device)[None, :].long() - trg = self.network.preprocess_target(trg) + trg = self.network.target_embedding(trg) logits = self.network.decoder(trg=trg, memory=memory, trg_mask=None) # Convert logits to probabilities. @@ -101,7 +98,7 @@ class VisionTransformerModel(Model): self.eval() if image.dtype == np.uint8: - # Converts an image with range [0, 255] with to Pytorch Tensor with range [0, 1]. + # 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. diff --git a/src/text_recognizer/networks/__init__.py b/src/text_recognizer/networks/__init__.py index 6d88768..2cc1137 100644 --- a/src/text_recognizer/networks/__init__.py +++ b/src/text_recognizer/networks/__init__.py @@ -1,25 +1,20 @@ """Network modules.""" from .cnn_transformer import CNNTransformer -from .cnn_transformer_encoder import CNNTransformerEncoder from .crnn import ConvolutionalRecurrentNetwork from .ctc import greedy_decoder from .densenet import DenseNet from .lenet import LeNet -from .loss import EmbeddingLoss from .mlp import MLP from .residual_network import ResidualNetwork, ResidualNetworkEncoder from .sparse_mlp import SparseMLP from .transformer import Transformer from .util import sliding_window -from .vision_transformer import VisionTransformer from .wide_resnet import WideResidualNetwork __all__ = [ "CNNTransformer", - "CNNTransformerEncoder", "ConvolutionalRecurrentNetwork", "DenseNet", - "EmbeddingLoss", "greedy_decoder", "MLP", "LeNet", @@ -28,6 +23,5 @@ __all__ = [ "sliding_window", "Transformer", "SparseMLP", - "VisionTransformer", "WideResidualNetwork", ] diff --git a/src/text_recognizer/networks/cnn_transformer.py b/src/text_recognizer/networks/cnn_transformer.py index 3da2c9f..16c7a41 100644 --- a/src/text_recognizer/networks/cnn_transformer.py +++ b/src/text_recognizer/networks/cnn_transformer.py @@ -1,4 +1,4 @@ -"""A DETR style transfomers but for text recognition.""" +"""A CNN-Transformer for image to text recognition.""" from typing import Dict, Optional, Tuple from einops import rearrange @@ -11,7 +11,7 @@ from text_recognizer.networks.util import configure_backbone class CNNTransformer(nn.Module): - """CNN+Transfomer for image to sequence prediction, sort of based on the ideas from DETR.""" + """CNN+Transfomer for image to sequence prediction.""" def __init__( self, @@ -25,22 +25,14 @@ class CNNTransformer(nn.Module): dropout_rate: float, trg_pad_index: int, backbone: str, - out_channels: int, - max_len: int, backbone_args: Optional[Dict] = None, activation: str = "gelu", ) -> None: super().__init__() self.trg_pad_index = trg_pad_index - self.backbone = configure_backbone(backbone, backbone_args) self.character_embedding = nn.Embedding(vocab_size, hidden_dim) - - # self.conv = nn.Conv2d(out_channels, max_len, kernel_size=1) - self.position_encoding = PositionalEncoding(hidden_dim, dropout_rate) - self.row_embed = nn.Parameter(torch.rand(max_len, max_len // 2)) - self.col_embed = nn.Parameter(torch.rand(max_len, max_len // 2)) self.adaptive_pool = ( nn.AdaptiveAvgPool2d((adaptive_pool_dim)) if adaptive_pool_dim else None @@ -78,8 +70,12 @@ class CNNTransformer(nn.Module): self.transformer.decoder(trg=trg, memory=memory, trg_mask=trg_mask) ) - def preprocess_input(self, src: Tensor) -> Tensor: - """Encodes src with a backbone network and a positional encoding. + def extract_image_features(self, src: Tensor) -> Tensor: + """Extracts image features with a backbone neural network. + + It seem like the winning idea was to swap channels and width dimension and collapse + the height dimension. The transformer is learning like a baby with this implementation!!! :D + Ohhhh, the joy I am experiencing right now!! Bring in the beers! :D :D :D Args: src (Tensor): Input tensor. @@ -88,29 +84,19 @@ class CNNTransformer(nn.Module): Tensor: A input src to the transformer. """ - # If batch dimenstion is missing, it needs to be added. + # If batch dimension is missing, it needs to be added. if len(src.shape) < 4: src = src[(None,) * (4 - len(src.shape))] src = self.backbone(src) - # src = self.conv(src) + src = rearrange(src, "b c h w -> b w c h") if self.adaptive_pool is not None: src = self.adaptive_pool(src) - H, W = src.shape[-2:] - src = rearrange(src, "b t h w -> b t (h w)") - - # construct positional encodings - pos = torch.cat( - [ - self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1), - self.row_embed[:H].unsqueeze(1).repeat(1, W, 1), - ], - dim=-1, - ).unsqueeze(0) - pos = rearrange(pos, "b h w l -> b l (h w)") - src = pos + 0.1 * src + src = src.squeeze(3) + src = self.position_encoding(src) + return src - def preprocess_target(self, trg: Tensor) -> Tuple[Tensor, Tensor]: + def target_embedding(self, trg: Tensor) -> Tuple[Tensor, Tensor]: """Encodes target tensor with embedding and postion. Args: @@ -126,9 +112,9 @@ class CNNTransformer(nn.Module): def forward(self, x: Tensor, trg: Optional[Tensor] = None) -> Tensor: """Forward pass with CNN transfomer.""" - h = self.preprocess_input(x) + h = self.extract_image_features(x) trg_mask = self._create_trg_mask(trg) - trg = self.preprocess_target(trg) + trg = self.target_embedding(trg) out = self.transformer(h, trg, trg_mask=trg_mask) logits = self.head(out) diff --git a/src/text_recognizer/networks/cnn_transformer_encoder.py b/src/text_recognizer/networks/cnn_transformer_encoder.py deleted file mode 100644 index 93626bf..0000000 --- a/src/text_recognizer/networks/cnn_transformer_encoder.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Network with a CNN backend and a transformer encoder head.""" -from typing import Dict - -from einops import rearrange -import torch -from torch import nn -from torch import Tensor - -from text_recognizer.networks.transformer import PositionalEncoding -from text_recognizer.networks.util import configure_backbone - - -class CNNTransformerEncoder(nn.Module): - """A CNN backbone with Transformer Encoder frontend for sequence prediction.""" - - def __init__( - self, - backbone: str, - backbone_args: Dict, - mlp_dim: int, - d_model: int, - nhead: int = 8, - dropout_rate: float = 0.1, - activation: str = "relu", - num_layers: int = 6, - num_classes: int = 80, - num_channels: int = 256, - max_len: int = 97, - ) -> None: - super().__init__() - self.d_model = d_model - self.nhead = nhead - self.dropout_rate = dropout_rate - self.activation = activation - self.num_layers = num_layers - - self.backbone = configure_backbone(backbone, backbone_args) - self.position_encoding = PositionalEncoding(d_model, dropout_rate) - self.encoder = self._configure_encoder() - - self.conv = nn.Conv2d(num_channels, max_len, kernel_size=1) - - self.mlp = nn.Linear(mlp_dim, d_model) - - self.head = nn.Linear(d_model, num_classes) - - def _configure_encoder(self) -> nn.TransformerEncoder: - encoder_layer = nn.TransformerEncoderLayer( - d_model=self.d_model, - nhead=self.nhead, - dropout=self.dropout_rate, - activation=self.activation, - ) - norm = nn.LayerNorm(self.d_model) - return nn.TransformerEncoder( - encoder_layer=encoder_layer, num_layers=self.num_layers, norm=norm - ) - - def forward(self, x: Tensor, targets: Tensor = None) -> Tensor: - """Forward pass through the network.""" - if len(x.shape) < 4: - x = x[(None,) * (4 - len(x.shape))] - - x = self.conv(self.backbone(x)) - x = rearrange(x, "b c h w -> b c (h w)") - x = self.mlp(x) - x = self.position_encoding(x) - x = rearrange(x, "b c h-> c b h") - x = self.encoder(x) - x = rearrange(x, "c b h-> b c h") - logits = self.head(x) - - return logits diff --git a/src/text_recognizer/networks/loss/__init__.py b/src/text_recognizer/networks/loss/__init__.py new file mode 100644 index 0000000..b489264 --- /dev/null +++ b/src/text_recognizer/networks/loss/__init__.py @@ -0,0 +1,2 @@ +"""Loss module.""" +from .loss import EmbeddingLoss, LabelSmoothingCrossEntropy diff --git a/src/text_recognizer/networks/loss.py b/src/text_recognizer/networks/loss/loss.py index cf9fa0d..cf9fa0d 100644 --- a/src/text_recognizer/networks/loss.py +++ b/src/text_recognizer/networks/loss/loss.py diff --git a/src/text_recognizer/networks/neural_machine_reader.py b/src/text_recognizer/networks/neural_machine_reader.py index 540a7d2..7f8c49b 100644 --- a/src/text_recognizer/networks/neural_machine_reader.py +++ b/src/text_recognizer/networks/neural_machine_reader.py @@ -1,180 +1,201 @@ -from typing import Dict, Optional, Tuple +"""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 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 +# from text_recognizer.networks.util import configure_backbone -class Encoder(nn.Module): +# 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) - 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) +# # 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, x: Tensor) -> Tuple[Tensor, Tensor]: - """Encodes a sequence of tensors with a bidirectional GRU. +# def forward(self, hidden_state: Tensor, encoder_outputs: Tensor) -> Tensor: +# """Short summary. - Args: - x (Tensor): A input sequence. +# 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)`. +# 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)) - where T is the sequence length, N is the batch size, E is the - embedding/encoder dimension, and D is the decoder dimension. +# a = self.attention(hidden_state, encoder_outputs) + +# weighted = torch.bmm(a, encoder_outputs) - 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, patch_size: Tuple[int, int] = (28, 28), - stride: Tuple[int, int] = (1, 14), dropout_rate: float = 0.1, teacher_forcing_ratio: float = 0.5) -> None: - super().__init__() - self.patch_size = patch_size - self.stride = stride - self.sliding_window = self._configure_sliding_window() - - self.backbone = - 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 _configure_sliding_window(self) -> nn.Sequential: - return nn.Sequential( - nn.Unfold(kernel_size=self.patch_size, stride=self.stride), - Rearrange( - "b (c h w) t -> b t c h w", - h=self.patch_size[0], - w=self.patch_size[1], - c=1, - ), - ) - - def forward(self, x: Tensor, trg: Tensor) -> Tensor: - #x = [batch size, height, width] - #trg = [trg len, batch size] - - # Create image patches with a sliding window kernel. - x = self.sliding_window(x) - - # Rearrange from a sequence of patches for feedforward network. - b, t = x.shape[:2] - x = rearrange(x, "b t c h w -> (b t) c h w", b=b, t=t) - - x = self.backbone(x) - x = rearrange(x, "(b t) h -> t b h", b=b, t=t) +# # 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/stn.py b/src/text_recognizer/networks/stn.py index b031128..e9d216f 100644 --- a/src/text_recognizer/networks/stn.py +++ b/src/text_recognizer/networks/stn.py @@ -13,7 +13,7 @@ class SpatialTransformerNetwork(nn.Module): Network that learns how to perform spatial transformations on the input image in order to enhance the geometric invariance of the model. - # TODO: add arguements to make it more general. + # TODO: add arguments to make it more general. """ diff --git a/src/text_recognizer/networks/util.py b/src/text_recognizer/networks/util.py index b31e640..e2d7955 100644 --- a/src/text_recognizer/networks/util.py +++ b/src/text_recognizer/networks/util.py @@ -24,7 +24,7 @@ def sliding_window( """ unfold = nn.Unfold(kernel_size=patch_size, stride=stride) - # Preform the slidning window, unsqueeze as the channel dimesion is lost. + # Preform the sliding window, unsqueeze as the channel dimesion is lost. c = images.shape[1] patches = unfold(images) patches = rearrange( diff --git a/src/text_recognizer/networks/vision_transformer.py b/src/text_recognizer/networks/vision_transformer.py deleted file mode 100644 index f227954..0000000 --- a/src/text_recognizer/networks/vision_transformer.py +++ /dev/null @@ -1,159 +0,0 @@ -"""VisionTransformer module. - -Splits each image into patches and feeds them to a transformer. - -""" - -from typing import Dict, Optional, Tuple, Type - -from einops import rearrange, reduce -from einops.layers.torch import Rearrange -from loguru import logger -import torch -from torch import nn -from torch import Tensor - -from text_recognizer.networks.transformer import PositionalEncoding, Transformer -from text_recognizer.networks.util import configure_backbone - - -class VisionTransformer(nn.Module): - """Linear projection+Transfomer for image to sequence prediction, sort of based on the ideas from ViT.""" - - def __init__( - self, - num_encoder_layers: int, - num_decoder_layers: int, - hidden_dim: int, - vocab_size: int, - num_heads: int, - max_len: int, - expansion_dim: int, - dropout_rate: float, - trg_pad_index: int, - mlp_dim: Optional[int] = None, - patch_size: Tuple[int, int] = (28, 28), - stride: Tuple[int, int] = (1, 14), - activation: str = "gelu", - backbone: Optional[str] = None, - backbone_args: Optional[Dict] = None, - ) -> None: - super().__init__() - - self.patch_size = patch_size - self.stride = stride - self.trg_pad_index = trg_pad_index - self.slidning_window = self._configure_sliding_window() - self.character_embedding = nn.Embedding(vocab_size, hidden_dim) - self.position_encoding = PositionalEncoding(hidden_dim, dropout_rate, max_len) - self.mlp_dim = mlp_dim - - self.use_backbone = False - if backbone is None: - self.linear_projection = nn.Linear( - self.patch_size[0] * self.patch_size[1], hidden_dim - ) - else: - self.backbone = configure_backbone(backbone, backbone_args) - if mlp_dim: - self.mlp = nn.Linear(mlp_dim, hidden_dim) - self.use_backbone = True - - self.transformer = Transformer( - num_encoder_layers, - num_decoder_layers, - hidden_dim, - num_heads, - expansion_dim, - dropout_rate, - activation, - ) - - self.head = nn.Sequential(nn.Linear(hidden_dim, vocab_size),) - - def _configure_sliding_window(self) -> nn.Sequential: - return nn.Sequential( - nn.Unfold(kernel_size=self.patch_size, stride=self.stride), - Rearrange( - "b (c h w) t -> b t c h w", - h=self.patch_size[0], - w=self.patch_size[1], - c=1, - ), - ) - - def _create_trg_mask(self, trg: Tensor) -> Tensor: - # Move this outside the transformer. - trg_pad_mask = (trg != self.trg_pad_index)[:, None, None] - trg_len = trg.shape[1] - trg_sub_mask = torch.tril( - torch.ones((trg_len, trg_len), device=trg.device) - ).bool() - trg_mask = trg_pad_mask & trg_sub_mask - return trg_mask - - def encoder(self, src: Tensor) -> Tensor: - """Forward pass with the encoder of the transformer.""" - return self.transformer.encoder(src) - - def decoder(self, trg: Tensor, memory: Tensor, trg_mask: Tensor) -> Tensor: - """Forward pass with the decoder of the transformer + classification head.""" - return self.head( - self.transformer.decoder(trg=trg, memory=memory, trg_mask=trg_mask) - ) - - def _backbone(self, x: Tensor) -> Tensor: - b, t = x.shape[:2] - if self.use_backbone: - x = rearrange(x, "b t c h w -> (b t) c h w", b=b, t=t) - x = self.backbone(x) - if self.mlp_dim: - x = rearrange(x, "(b t) c h w -> b t (c h w)", b=b, t=t) - x = self.mlp(x) - else: - x = rearrange(x, "(b t) h -> b t h", b=b, t=t) - else: - x = rearrange(x, "b t c h w -> b t (c h w)", b=b, t=t) - x = self.linear_projection(x) - return x - - def preprocess_input(self, src: Tensor) -> Tensor: - """Encodes src with a backbone network and a positional encoding. - - Args: - src (Tensor): Input tensor. - - Returns: - Tensor: A input src to the transformer. - - """ - # If batch dimenstion is missing, it needs to be added. - if len(src.shape) < 4: - src = src[(None,) * (4 - len(src.shape))] - src = self.slidning_window(src) # .squeeze(-2) - src = self._backbone(src) - src = self.position_encoding(src) - return src - - def preprocess_target(self, trg: Tensor) -> Tuple[Tensor, Tensor]: - """Encodes target tensor with embedding and postion. - - Args: - trg (Tensor): Target tensor. - - Returns: - Tuple[Tensor, Tensor]: Encoded target tensor and target mask. - - """ - trg_mask = self._create_trg_mask(trg) - trg = self.character_embedding(trg.long()) - trg = self.position_encoding(trg) - return trg, trg_mask - - def forward(self, x: Tensor, trg: Tensor) -> Tensor: - """Forward pass with vision transfomer.""" - src = self.preprocess_input(x) - trg, trg_mask = self.preprocess_target(trg) - out = self.transformer(src, trg, trg_mask=trg_mask) - logits = self.head(out) - return logits diff --git a/src/text_recognizer/networks/wide_resnet.py b/src/text_recognizer/networks/wide_resnet.py index aa79c12..28f3380 100644 --- a/src/text_recognizer/networks/wide_resnet.py +++ b/src/text_recognizer/networks/wide_resnet.py @@ -2,7 +2,7 @@ from functools import partial from typing import Callable, Dict, List, Optional, Type, Union -from einops.layers.torch import Rearrange, Reduce +from einops.layers.torch import Reduce import numpy as np import torch from torch import nn diff --git a/src/text_recognizer/weights/CharacterModel_EmnistDataset_DenseNet_weights.pt b/src/text_recognizer/weights/CharacterModel_EmnistDataset_DenseNet_weights.pt index 4c28b51..f2dfd84 100644 --- a/src/text_recognizer/weights/CharacterModel_EmnistDataset_DenseNet_weights.pt +++ b/src/text_recognizer/weights/CharacterModel_EmnistDataset_DenseNet_weights.pt @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:2a7fe47aec144ea79e50a145e2fbc29e0174092c9e702c0cd0c6ef105ca99760 -size 1273881 +oid sha256:8a69e5efedea70c4c5cb8ccdcc8cd480400f6c73e3313423f4dbbfe615644f0a +size 4500617 diff --git a/src/text_recognizer/weights/CharacterModel_EmnistDataset_WideResidualNetwork_weights.pt b/src/text_recognizer/weights/CharacterModel_EmnistDataset_WideResidualNetwork_weights.pt index cd3e37a..e1add8d 100644 --- a/src/text_recognizer/weights/CharacterModel_EmnistDataset_WideResidualNetwork_weights.pt +++ b/src/text_recognizer/weights/CharacterModel_EmnistDataset_WideResidualNetwork_weights.pt @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4d54658ffc40698149d0db1f9c0b4e446809f6dc237d9a384c54cb2e69219af2 -size 14953410 +oid sha256:68dd5c98eedc8753546f88b4e6fd5fc38725dc0079b837c30fb3d48069ec412b +size 15002754 |