IAM datasets implemented.

5 files changed, 417 insertions, 124 deletions

diff --git a/src/text_recognizer/models/__init__.py b/src/text_recognizer/models/__init__.py
index ff10a07..a3cfc15 100644
--- a/src/text_recognizer/models/__init__.py
+++ b/src/text_recognizer/models/__init__.py
@@ -1,6 +1,7 @@
 """Model modules."""
 from .base import Model
 from .character_model import CharacterModel
-from .metrics import accuracy
+from .line_ctc_model import LineCTCModel
+from .metrics import accuracy, cer, wer
 
-__all__ = ["Model", "CharacterModel", "accuracy"]
+__all__ = ["Model", "cer", "CharacterModel", "LineCTCModel", "accuracy", "wer"]
diff --git a/src/text_recognizer/models/base.py b/src/text_recognizer/models/base.py
index 3a84a11..153e19a 100644
--- a/src/text_recognizer/models/base.py
+++ b/src/text_recognizer/models/base.py
@@ -2,6 +2,7 @@
 
 from abc import ABC, abstractmethod
 from glob import glob
+import importlib
 from pathlib import Path
 import re
 import shutil
@@ -10,9 +11,12 @@ from typing import Callable, Dict, Optional, Tuple, Type
 from loguru import logger
 import torch
 from torch import nn
+from torch import Tensor
+from torch.optim.swa_utils import AveragedModel, SWALR
+from torch.utils.data import DataLoader, Dataset, random_split
 from torchsummary import summary
 
-from text_recognizer.datasets import EmnistMapper, fetch_data_loaders
+from text_recognizer.datasets import EmnistMapper
 
 WEIGHT_DIRNAME = Path(__file__).parents[1].resolve() / "weights"
 
@@ -23,8 +27,9 @@ class Model(ABC):
     def __init__(
         self,
         network_fn: Type[nn.Module],
+        dataset: Type[Dataset],
         network_args: Optional[Dict] = None,
-        data_loader_args: Optional[Dict] = None,
+        dataset_args: Optional[Dict] = None,
         metrics: Optional[Dict] = None,
         criterion: Optional[Callable] = None,
         criterion_args: Optional[Dict] = None,
@@ -32,14 +37,16 @@ class Model(ABC):
         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:
         """Base class, to be inherited by model for specific type of data.
 
         Args:
             network_fn (Type[nn.Module]): The PyTorch network.
+            dataset (Type[Dataset]): A dataset class.
             network_args (Optional[Dict]): Arguments for the network. Defaults to None.
-            data_loader_args (Optional[Dict]):  Arguments for the DataLoader.
+            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.
                 Defaults to None.
@@ -49,107 +56,181 @@ class Model(ABC):
             lr_scheduler (Optional[Callable]): A PyTorch learning rate scheduler. Defaults to None.
             lr_scheduler_args (Optional[Dict]): Dict of arguments for learning rate scheduler. Defaults to
                 None.
+            swa_args (Optional[Dict]): Dict of arguments for stochastic weight averaging. Defaults to
+                None.
             device (Optional[str]): Name of the device to train on. Defaults to None.
 
         """
+        # Has to be set in subclass.
+        self._mapper = None
 
-        # Configure data loaders and dataset info.
-        dataset_name, self._data_loaders, self._mapper = self._configure_data_loader(
-            data_loader_args
-        )
-        self._input_shape = self._mapper.input_shape
+        # Placeholder.
+        self._input_shape = None
+
+        self.dataset = dataset
+        self.dataset_args = dataset_args
+
+        # Placeholders for datasets.
+        self.train_dataset = None
+        self.val_dataset = None
+        self.test_dataset = None
 
-        self._name = f"{self.__class__.__name__}_{dataset_name}_{network_fn.__name__}"
+        # Stochastic Weight Averaging placeholders.
+        self.swa_args = swa_args
+        self._swa_start = None
+        self._swa_scheduler = None
+        self._swa_network = None
 
-        if metrics is not None:
-            self._metrics = metrics
+        # Experiment directory.
+        self.model_dir = None
+
+        # Flag for configured model.
+        self.is_configured = False
+        self.data_prepared = False
+
+        # Flag for stopping training.
+        self.stop_training = False
+
+        self._name = (
+            f"{self.__class__.__name__}_{dataset.__name__}_{network_fn.__name__}"
+        )
+
+        self._metrics = metrics if metrics is not None else None
 
         # Set the device.
-        if device is None:
-            self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        else:
-            self._device = device
+        self._device = (
+            torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if device is None
+            else device
+        )
 
         # Configure network.
-        self._network, self._network_args = self._configure_network(
-            network_fn, network_args
-        )
+        self._network = None
+        self._network_args = network_args
+        self._configure_network(network_fn)
 
-        # To device.
-        self._network.to(self._device)
+        # Place network on device (GPU).
+        self.to_device()
+
+        # Loss and Optimizer placeholders for before loading.
+        self._criterion = criterion
+        self.criterion_args = criterion_args
+
+        self._optimizer = optimizer
+        self.optimizer_args = optimizer_args
+
+        self._lr_scheduler = lr_scheduler
+        self.lr_scheduler_args = lr_scheduler_args
+
+    def configure_model(self) -> None:
+        """Configures criterion and optimizers."""
+        if not self.is_configured:
+            self._configure_criterion()
+            self._configure_optimizers()
+
+            # Prints a summary of the network in terminal.
+            self.summary()
+
+            # Set this flag to true to prevent the model from configuring again.
+            self.is_configured = True
+
+    def prepare_data(self) -> None:
+        """Prepare data for training."""
+        # TODO add downloading.
+        if not self.data_prepared:
+            # Load train dataset.
+            train_dataset = self.dataset(train=True, **self.dataset_args["args"])
+
+            # Set input shape.
+            self._input_shape = train_dataset.input_shape
+
+            # Split train dataset into a training and validation partition.
+            dataset_len = len(train_dataset)
+            train_len = int(
+                self.dataset_args["train_args"]["train_fraction"] * dataset_len
+            )
+            val_len = dataset_len - train_len
+            self.train_dataset, self.val_dataset = random_split(
+                train_dataset, lengths=[train_len, val_len]
+            )
+
+            # Load test dataset.
+            self.test_dataset = self.dataset(train=False, **self.dataset_args["args"])
+
+            # Set the flag to true to disable ability to load data agian.
+            self.data_prepared = True
 
-        # Configure training objects.
-        self._criterion = self._configure_criterion(criterion, criterion_args)
-        self._optimizer, self._lr_scheduler = self._configure_optimizers(
-            optimizer, optimizer_args, lr_scheduler, lr_scheduler_args
+    def train_dataloader(self) -> DataLoader:
+        """Returns data loader for training set."""
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.dataset_args["train_args"]["batch_size"],
+            num_workers=self.dataset_args["train_args"]["num_workers"],
+            shuffle=True,
+            pin_memory=True,
         )
 
-        # Experiment directory.
-        self.model_dir = None
+    def val_dataloader(self) -> DataLoader:
+        """Returns data loader for validation set."""
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.dataset_args["train_args"]["batch_size"],
+            num_workers=self.dataset_args["train_args"]["num_workers"],
+            shuffle=True,
+            pin_memory=True,
+        )
 
-        # Flag for stopping training.
-        self.stop_training = False
+    def test_dataloader(self) -> DataLoader:
+        """Returns data loader for test set."""
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.dataset_args["train_args"]["batch_size"],
+            num_workers=self.dataset_args["train_args"]["num_workers"],
+            shuffle=False,
+            pin_memory=True,
+        )
 
-    def _configure_data_loader(
-        self, data_loader_args: Optional[Dict]
-    ) -> Tuple[str, Dict, EmnistMapper]:
-        """Loads data loader, dataset name, and dataset mapper."""
-        if data_loader_args is not None:
-            data_loaders = fetch_data_loaders(**data_loader_args)
-            dataset = list(data_loaders.values())[0].dataset
-            dataset_name = dataset.__name__
-            mapper = dataset.mapper
-        else:
-            self._mapper = EmnistMapper()
-            dataset_name = "*"
-            data_loaders = None
-        return dataset_name, data_loaders, mapper
-
-    def _configure_network(
-        self, network_fn: Type[nn.Module], network_args: Optional[Dict]
-    ) -> Tuple[Type[nn.Module], Dict]:
+    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 network_args is None:
-            network, network_args = self.load_weights(network_fn)
+        if self._network_args is None:
+            self.load_weights(network_fn)
         else:
-            network = network_fn(**network_args)
-        return network, network_args
+            self._network = network_fn(**self._network_args)
 
-    def _configure_criterion(
-        self, criterion: Optional[Callable], criterion_args: Optional[Dict]
-    ) -> Optional[Callable]:
+    def _configure_criterion(self) -> None:
         """Loads the criterion."""
-        if criterion is not None:
-            _criterion = criterion(**criterion_args)
-        else:
-            _criterion = None
-        return _criterion
+        self._criterion = (
+            self._criterion(**self.criterion_args)
+            if self._criterion is not None
+            else None
+        )
 
-    def _configure_optimizers(
-        self,
-        optimizer: Optional[Callable],
-        optimizer_args: Optional[Dict],
-        lr_scheduler: Optional[Callable],
-        lr_scheduler_args: Optional[Dict],
-    ) -> Tuple[Optional[Callable], Optional[Callable]]:
+    def _configure_optimizers(self,) -> None:
         """Loads the optimizers."""
-        if optimizer is not None:
-            _optimizer = optimizer(self._network.parameters(), **optimizer_args)
+        if self._optimizer is not None:
+            self._optimizer = self._optimizer(
+                self._network.parameters(), **self.optimizer_args
+            )
         else:
-            _optimizer = None
+            self._optimizer = None
 
-        if _optimizer and lr_scheduler is not None:
-            if "OneCycleLR" in str(lr_scheduler):
-                lr_scheduler_args["steps_per_epoch"] = len(self._data_loaders["train"])
-            _lr_scheduler = lr_scheduler(_optimizer, **lr_scheduler_args)
+        if self._optimizer and self._lr_scheduler is not None:
+            if "OneCycleLR" in str(self._lr_scheduler):
+                self.lr_scheduler_args["steps_per_epoch"] = len(self.train_dataloader())
+            self._lr_scheduler = self._lr_scheduler(
+                self._optimizer, **self.lr_scheduler_args
+            )
         else:
-            _lr_scheduler = None
+            self._lr_scheduler = None
 
-        return _optimizer, _lr_scheduler
+        if self.swa_args is not None:
+            self._swa_start = self.swa_args["start"]
+            self._swa_scheduler = SWALR(self._optimizer, swa_lr=self.swa_args["lr"])
+            self._swa_network = AveragedModel(self._network).to(self.device)
 
     @property
-    def __name__(self) -> str:
+    def name(self) -> str:
         """Returns the name of the model."""
         return self._name
 
@@ -159,7 +240,7 @@ class Model(ABC):
         return self._input_shape
 
     @property
-    def mapper(self) -> Dict:
+    def mapper(self) -> EmnistMapper:
         """Returns the mapper that maps between ints and chars."""
         return self._mapper
 
@@ -202,13 +283,24 @@ class Model(ABC):
         return self._lr_scheduler
 
     @property
-    def data_loaders(self) -> Optional[Dict]:
-        """Dataloaders."""
-        return self._data_loaders
+    def swa_scheduler(self) -> Optional[Callable]:
+        """Returns the stochastic weight averaging scheduler."""
+        return self._swa_scheduler
+
+    @property
+    def swa_start(self) -> Optional[Callable]:
+        """Returns the start epoch of stochastic weight averaging."""
+        return self._swa_start
 
     @property
-    def network(self) -> nn.Module:
+    def swa_network(self) -> Optional[Callable]:
+        """Returns the stochastic weight averaging network."""
+        return self._swa_network
+
+    @property
+    def network(self) -> Type[nn.Module]:
         """Neural network."""
+        # Returns the SWA network if available.
         return self._network
 
     @property
@@ -217,15 +309,27 @@ class Model(ABC):
         WEIGHT_DIRNAME.mkdir(parents=True, exist_ok=True)
         return str(WEIGHT_DIRNAME / f"{self._name}_weights.pt")
 
-    def summary(self) -> None:
+    def loss_fn(self, output: Tensor, targets: Tensor) -> Tensor:
+        """Compute the loss."""
+        return self.criterion(output, targets)
+
+    def summary(
+        self, input_shape: Optional[Tuple[int, int, int]] = None, depth: int = 5
+    ) -> None:
         """Prints a summary of the network architecture."""
-        device = re.sub("[^A-Za-z]+", "", self.device)
-        if self._input_shape is not None:
+
+        if input_shape is not None:
+            summary(self._network, input_shape, depth=depth, device=self.device)
+        elif self._input_shape is not None:
             input_shape = (1,) + tuple(self._input_shape)
-            summary(self._network, input_shape, device=device)
+            summary(self._network, input_shape, depth=depth, device=self.device)
         else:
             logger.warning("Could not print summary as input shape is not set.")
 
+    def to_device(self) -> None:
+        """Places the network on the device (GPU)."""
+        self._network.to(self._device)
+
     def _get_state_dict(self) -> Dict:
         """Get the state dict of the model."""
         state = {"model_state": self._network.state_dict()}
@@ -236,69 +340,67 @@ class Model(ABC):
         if self._lr_scheduler is not None:
             state["scheduler_state"] = self._lr_scheduler.state_dict()
 
+        if self._swa_network is not None:
+            state["swa_network"] = self._swa_network.state_dict()
+
         return state
 
-    def load_checkpoint(self, path: Path) -> int:
+    def load_from_checkpoint(self, checkpoint_path: Path) -> None:
         """Load a previously saved checkpoint.
 
         Args:
-            path (Path): Path to the experiment with the checkpoint.
-
-        Returns:
-            epoch (int): The last epoch when the checkpoint was created.
+            checkpoint_path (Path): Path to the experiment with the checkpoint.
 
         """
         logger.debug("Loading checkpoint...")
-        if not path.exists():
-            logger.debug("File does not exist {str(path)}")
+        if not checkpoint_path.exists():
+            logger.debug("File does not exist {str(checkpoint_path)}")
 
-        checkpoint = torch.load(str(path))
+        checkpoint = torch.load(str(checkpoint_path))
         self._network.load_state_dict(checkpoint["model_state"])
 
         if self._optimizer is not None:
             self._optimizer.load_state_dict(checkpoint["optimizer_state"])
 
-        # Does not work when loadning from previous checkpoint and trying to train beyond the last max epochs.
-        # if self._lr_scheduler is not None:
-        #     self._lr_scheduler.load_state_dict(checkpoint["scheduler_state"])
-
-        epoch = checkpoint["epoch"]
+        if self._lr_scheduler is not None:
+            # Does not work when loadning from previous checkpoint and trying to train beyond the last max epochs
+            # with OneCycleLR.
+            if self._lr_scheduler.__class__.__name__ != "OneCycleLR":
+                self._lr_scheduler.load_state_dict(checkpoint["scheduler_state"])
 
-        return epoch
+        if self._swa_network is not None:
+            self._swa_network.load_state_dict(checkpoint["swa_network"])
 
-    def save_checkpoint(self, is_best: bool, epoch: int, val_metric: str) -> None:
+    def save_checkpoint(
+        self, checkpoint_path: Path, is_best: bool, epoch: int, val_metric: str
+    ) -> None:
         """Saves a checkpoint of the model.
 
         Args:
+            checkpoint_path (Path): Path to the experiment with the checkpoint.
             is_best (bool): If it is the currently best model.
             epoch (int): The epoch of the checkpoint.
             val_metric (str): Validation metric.
 
-        Raises:
-            ValueError: If the self.model_dir is not set.
-
         """
         state = self._get_state_dict()
         state["is_best"] = is_best
         state["epoch"] = epoch
         state["network_args"] = self._network_args
 
-        if self.model_dir is None:
-            raise ValueError("Experiment directory is not set.")
-
-        self.model_dir.mkdir(parents=True, exist_ok=True)
+        checkpoint_path.mkdir(parents=True, exist_ok=True)
 
         logger.debug("Saving checkpoint...")
-        filepath = str(self.model_dir / "last.pt")
+        filepath = str(checkpoint_path / "last.pt")
         torch.save(state, filepath)
 
         if is_best:
             logger.debug(
                 f"Found a new best {val_metric}. Saving best checkpoint and weights."
             )
-            shutil.copyfile(filepath, str(self.model_dir / "best.pt"))
+            shutil.copyfile(filepath, str(checkpoint_path / "best.pt"))
 
-    def load_weights(self, network_fn: Type[nn.Module]) -> Tuple[Type[nn.Module], Dict]:
+    def load_weights(self, network_fn: Type[nn.Module]) -> None:
         """Load the network weights."""
         logger.debug("Loading network with pretrained weights.")
         filename = glob(self.weights_filename)[0]
@@ -308,13 +410,16 @@ class Model(ABC):
             )
         # Loading state directory.
         state_dict = torch.load(filename, map_location=torch.device(self._device))
-        network_args = state_dict["network_args"]
+        self._network_args = state_dict["network_args"]
         weights = state_dict["model_state"]
 
         # Initializes the network with trained weights.
-        network = network_fn(**self._network_args)
-        network.load_state_dict(weights)
-        return network, network_args
+        self._network = network_fn(**self._network_args)
+        self._network.load_state_dict(weights)
+
+        if "swa_network" in state_dict:
+            self._swa_network = AveragedModel(self._network).to(self.device)
+            self._swa_network.load_state_dict(state_dict["swa_network"])
 
     def save_weights(self, path: Path) -> None:
         """Save the network weights."""
diff --git a/src/text_recognizer/models/character_model.py b/src/text_recognizer/models/character_model.py
index 0fd7afd..64ba693 100644
--- a/src/text_recognizer/models/character_model.py
+++ b/src/text_recognizer/models/character_model.py
@@ -4,8 +4,10 @@ from typing import Callable, Dict, Optional, Tuple, Type, Union
 import numpy as np
 import torch
 from torch import nn
+from torch.utils.data import Dataset
 from torchvision.transforms import ToTensor
 
+from text_recognizer.datasets import EmnistMapper
 from text_recognizer.models.base import Model
 
 
@@ -15,8 +17,9 @@ class CharacterModel(Model):
     def __init__(
         self,
         network_fn: Type[nn.Module],
+        dataset: Type[Dataset],
         network_args: Optional[Dict] = None,
-        data_loader_args: Optional[Dict] = None,
+        dataset_args: Optional[Dict] = None,
         metrics: Optional[Dict] = None,
         criterion: Optional[Callable] = None,
         criterion_args: Optional[Dict] = None,
@@ -24,14 +27,16 @@ class CharacterModel(Model):
         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:
         """Initializes the CharacterModel."""
 
         super().__init__(
             network_fn,
+            dataset,
             network_args,
-            data_loader_args,
+            dataset_args,
             metrics,
             criterion,
             criterion_args,
@@ -39,8 +44,11 @@ class CharacterModel(Model):
             optimizer_args,
             lr_scheduler,
             lr_scheduler_args,
+            swa_args,
             device,
         )
+        if self._mapper is None:
+            self._mapper = EmnistMapper()
         self.tensor_transform = ToTensor()
         self.softmax = nn.Softmax(dim=0)
 
@@ -67,9 +75,13 @@ class CharacterModel(Model):
 
         # Put the image tensor on the device the model weights are on.
         image = image.to(self.device)
-        logits = self.network(image)
+        logits = (
+            self.swa_network(image)
+            if self.swa_network is not None
+            else self.network(image)
+        )
 
-        prediction = self.softmax(logits.data.squeeze())
+        prediction = self.softmax(logits.squeeze(0))
 
         index = int(torch.argmax(prediction, dim=0))
         confidence_of_prediction = prediction[index]
diff --git a/src/text_recognizer/models/line_ctc_model.py b/src/text_recognizer/models/line_ctc_model.py
new file mode 100644
index 0000000..97308a7
--- /dev/null
+++ b/src/text_recognizer/models/line_ctc_model.py
@@ -0,0 +1,105 @@
+"""Defines the LineCTCModel class."""
+from typing import Callable, Dict, 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
+from text_recognizer.networks import greedy_decoder
+
+
+class LineCTCModel(Model):
+    """Model for predicting a sequence of characters from an image of a text line."""
+
+    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,
+        )
+        if self._mapper is None:
+            self._mapper = EmnistMapper()
+        self.tensor_transform = ToTensor()
+
+    def loss_fn(self, output: Tensor, targets: Tensor) -> Tensor:
+        """Computes the CTC loss.
+
+        Args:
+            output (Tensor): Model predictions.
+            targets (Tensor): Correct output sequence.
+
+        Returns:
+            Tensor: The CTC loss.
+
+        """
+        input_lengths = torch.full(
+            size=(output.shape[1],), fill_value=output.shape[0], dtype=torch.long,
+        )
+        target_lengths = torch.full(
+            size=(output.shape[1],), fill_value=targets.shape[1], dtype=torch.long,
+        )
+        return self.criterion(output, targets, input_lengths, target_lengths)
+
+    @torch.no_grad()
+    def predict_on_image(self, image: Union[np.ndarray, Tensor]) -> Tuple[str, float]:
+        """Predict on a single input."""
+        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)
+        log_probs = (
+            self.swa_network(image)
+            if self.swa_network is not None
+            else self.network(image)
+        )
+
+        raw_pred, _ = greedy_decoder(
+            predictions=log_probs,
+            character_mapper=self.mapper,
+            blank_label=79,
+            collapse_repeated=True,
+        )
+
+        log_probs, _ = log_probs.max(dim=2)
+
+        predicted_characters = "".join(raw_pred[0])
+        confidence_of_prediction = torch.exp(log_probs.sum()).item()
+
+        return predicted_characters, confidence_of_prediction
diff --git a/src/text_recognizer/models/metrics.py b/src/text_recognizer/models/metrics.py
index ac8d68e..6a26216 100644
--- a/src/text_recognizer/models/metrics.py
+++ b/src/text_recognizer/models/metrics.py
@@ -1,19 +1,89 @@
 """Utility functions for models."""
-
+import Levenshtein as Lev
 import torch
+from torch import Tensor
+
+from text_recognizer.networks import greedy_decoder
 
 
-def accuracy(outputs: torch.Tensor, labels: torch.Tensor) -> float:
+def accuracy(outputs: Tensor, labels: Tensor) -> float:
     """Computes the accuracy.
 
     Args:
-        outputs (torch.Tensor): The output from the network.
-        labels (torch.Tensor): Ground truth labels.
+        outputs (Tensor): The output from the network.
+        labels (Tensor): Ground truth labels.
 
     Returns:
         float: The accuracy for the batch.
 
     """
     _, predicted = torch.max(outputs.data, dim=1)
-    acc = (predicted == labels).sum().item() / labels.shape[0]
+    acc = (predicted == labels).sum().float() / labels.shape[0]
+    acc = acc.item()
     return acc
+
+
+def cer(outputs: Tensor, targets: Tensor) -> float:
+    """Computes the character error rate.
+
+    Args:
+        outputs (Tensor): The output from the network.
+        targets (Tensor): Ground truth labels.
+
+    Returns:
+        float: The cer for the batch.
+
+    """
+    target_lengths = torch.full(
+        size=(outputs.shape[1],), fill_value=targets.shape[1], dtype=torch.long,
+    )
+    decoded_predictions, decoded_targets = greedy_decoder(
+        outputs, targets, target_lengths
+    )
+
+    lev_dist = 0
+
+    for prediction, target in zip(decoded_predictions, decoded_targets):
+        prediction = "".join(prediction)
+        target = "".join(target)
+        prediction, target = (
+            prediction.replace(" ", ""),
+            target.replace(" ", ""),
+        )
+        lev_dist += Lev.distance(prediction, target)
+    return lev_dist / len(decoded_predictions)
+
+
+def wer(outputs: Tensor, targets: Tensor) -> float:
+    """Computes the Word error rate.
+
+    Args:
+        outputs (Tensor): The output from the network.
+        targets (Tensor): Ground truth labels.
+
+    Returns:
+        float: The wer for the batch.
+
+    """
+    target_lengths = torch.full(
+        size=(outputs.shape[1],), fill_value=targets.shape[1], dtype=torch.long,
+    )
+    decoded_predictions, decoded_targets = greedy_decoder(
+        outputs, targets, target_lengths
+    )
+
+    lev_dist = 0
+
+    for prediction, target in zip(decoded_predictions, decoded_targets):
+        prediction = "".join(prediction)
+        target = "".join(target)
+
+        b = set(prediction.split() + target.split())
+        word2char = dict(zip(b, range(len(b))))
+
+        w1 = [chr(word2char[w]) for w in prediction.split()]
+        w2 = [chr(word2char[w]) for w in target.split()]
+
+        lev_dist += Lev.distance("".join(w1), "".join(w2))
+
+    return lev_dist / len(decoded_predictions)