onnxtr 0.1.0__py3-none-any.whl

onnxtr/models/preprocessor/base.py

@@ -0,0 +1,118 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+import math
+from typing import Any, List, Tuple, Union
+
+import numpy as np
+
+from onnxtr.transforms import Normalize, Resize
+from onnxtr.utils.geometry import shape_translate
+from onnxtr.utils.multithreading import multithread_exec
+from onnxtr.utils.repr import NestedObject
+
+__all__ = ["PreProcessor"]
+
+
+class PreProcessor(NestedObject):
+    """Implements an abstract preprocessor object which performs casting, resizing, batching and normalization.
+
+    Args:
+    ----
+        output_size: expected size of each page in format (H, W)
+        batch_size: the size of page batches
+        mean: mean value of the training distribution by channel
+        std: standard deviation of the training distribution by channel
+    """
+
+    _children_names: List[str] = ["resize", "normalize"]
+
+    def __init__(
+        self,
+        output_size: Tuple[int, int],
+        batch_size: int,
+        mean: Tuple[float, float, float] = (0.5, 0.5, 0.5),
+        std: Tuple[float, float, float] = (1.0, 1.0, 1.0),
+        **kwargs: Any,
+    ) -> None:
+        self.batch_size = batch_size
+        self.resize = Resize(output_size, **kwargs)
+        self.normalize = Normalize(mean, std)
+
+    def batch_inputs(self, samples: List[np.ndarray]) -> List[np.ndarray]:
+        """Gather samples into batches for inference purposes
+
+        Args:
+        ----
+            samples: list of samples (tf.Tensor)
+
+        Returns:
+        -------
+            list of batched samples
+        """
+        num_batches = int(math.ceil(len(samples) / self.batch_size))
+        batches = [
+            np.stack(samples[idx * self.batch_size : min((idx + 1) * self.batch_size, len(samples))], axis=0)
+            for idx in range(int(num_batches))
+        ]
+
+        return batches
+
+    def sample_transforms(self, x: np.ndarray) -> np.ndarray:
+        if x.ndim != 3:
+            raise AssertionError("expected list of 3D Tensors")
+        if isinstance(x, np.ndarray):
+            if x.dtype not in (np.uint8, np.float32):
+                raise TypeError("unsupported data type for numpy.ndarray")
+        x = shape_translate(x, "HWC")
+        # Data type & 255 division
+        if x.dtype == np.uint8:
+            x = x.astype(np.float32) / 255.0
+        # Resizing
+        x = self.resize(x)
+
+        return x
+
+    def __call__(self, x: Union[np.ndarray, List[np.ndarray]]) -> List[np.ndarray]:
+        """Prepare document data for model forwarding
+
+        Args:
+        ----
+            x: list of images (np.array) or tensors (already resized and batched)
+
+        Returns:
+        -------
+            list of page batches
+        """
+        # Input type check
+        if isinstance(x, np.ndarray):
+            if x.ndim != 4:
+                raise AssertionError("expected 4D Tensor")
+            if isinstance(x, np.ndarray):
+                if x.dtype not in (np.uint8, np.float32):
+                    raise TypeError("unsupported data type for numpy.ndarray")
+            x = shape_translate(x, "BHWC")
+
+            # Data type & 255 division
+            if x.dtype == np.uint8:
+                x = x.astype(np.float32) / 255.0
+            # Resizing
+            if (x.shape[1], x.shape[2]) != self.resize.output_size:
+                x = np.array([self.resize(sample) for sample in x])
+
+            batches = [x]
+
+        elif isinstance(x, list) and all(isinstance(sample, np.ndarray) for sample in x):
+            # Sample transform (to tensor, resize)
+            samples = list(multithread_exec(self.sample_transforms, x))
+            # Batching
+            batches = self.batch_inputs(samples)
+        else:
+            raise TypeError(f"invalid input type: {type(x)}")
+
+        # Batch transforms (normalize)
+        batches = list(multithread_exec(self.normalize, batches))
+
+        return batches

onnxtr/models/recognition/__init__.py

@@ -0,0 +1,2 @@
+from .models import *
+from .zoo import *

onnxtr/models/recognition/core.py

@@ -0,0 +1,28 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+
+from onnxtr.utils.repr import NestedObject
+
+__all__ = ["RecognitionPostProcessor"]
+
+
+class RecognitionPostProcessor(NestedObject):
+    """Abstract class to postprocess the raw output of the model
+
+    Args:
+    ----
+        vocab: string containing the ordered sequence of supported characters
+    """
+
+    def __init__(
+        self,
+        vocab: str,
+    ) -> None:
+        self.vocab = vocab
+        self._embedding = list(self.vocab) + ["<eos>"]
+
+    def extra_repr(self) -> str:
+        return f"vocab_size={len(self.vocab)}"

onnxtr/models/recognition/models/__init__.py

@@ -0,0 +1,5 @@
+from .crnn import *
+from .sar import *
+from .master import *
+from .vitstr import *
+from .parseq import *

onnxtr/models/recognition/models/crnn.py

@@ -0,0 +1,226 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from copy import deepcopy
+from itertools import groupby
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+from scipy.special import softmax
+
+from onnxtr.utils import VOCABS
+
+from ...engine import Engine
+from ..core import RecognitionPostProcessor
+
+__all__ = ["CRNN", "crnn_vgg16_bn", "crnn_mobilenet_v3_small", "crnn_mobilenet_v3_large"]
+
+default_cfgs: Dict[str, Dict[str, Any]] = {
+    "crnn_vgg16_bn": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 32, 128),
+        "vocab": VOCABS["legacy_french"],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/crnn_vgg16_bn-662979cc.onnx",
+    },
+    "crnn_mobilenet_v3_small": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 32, 128),
+        "vocab": VOCABS["french"],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/crnn_mobilenet_v3_small-bded4d49.onnx",
+    },
+    "crnn_mobilenet_v3_large": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 32, 128),
+        "vocab": VOCABS["french"],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/crnn_mobilenet_v3_large-d42e8185.onnx",
+    },
+}
+
+
+class CRNNPostProcessor(RecognitionPostProcessor):
+    """Postprocess raw prediction of the model (logits) to a list of words using CTC decoding
+
+    Args:
+    ----
+        vocab: string containing the ordered sequence of supported characters
+    """
+
+    def __init__(self, vocab):
+        self.vocab = vocab
+
+    def decode_sequence(self, sequence, vocab):
+        return "".join([vocab[int(char)] for char in sequence])
+
+    def ctc_best_path(
+        self,
+        logits,
+        vocab,
+        blank=0,
+    ):
+        """Implements best path decoding as shown by Graves (Dissertation, p63), highly inspired from
+        <https://github.com/githubharald/CTCDecoder>`_.
+
+        Args:
+        ----
+            logits: model output, shape: N x T x C
+            vocab: vocabulary to use
+            blank: index of blank label
+
+        Returns:
+        -------
+            A list of tuples: (word, confidence)
+        """
+        # Gather the most confident characters, and assign the smallest conf among those to the sequence prob
+        probs = softmax(logits, axis=-1).max(axis=-1).min(axis=1)
+
+        # collapse best path (using itertools.groupby), map to chars, join char list to string
+        words = [
+            self.decode_sequence([k for k, _ in groupby(seq.tolist()) if k != blank], vocab)
+            for seq in np.argmax(logits, axis=-1)
+        ]
+
+        return list(zip(words, probs.astype(float).tolist()))
+
+    def __call__(self, logits):
+        """Performs decoding of raw output with CTC and decoding of CTC predictions
+        with label_to_idx mapping dictionnary
+
+        Args:
+        ----
+            logits: raw output of the model, shape (N, C + 1, seq_len)
+
+        Returns:
+        -------
+            A tuple of 2 lists: a list of str (words) and a list of float (probs)
+
+        """
+        # Decode CTC
+        return self.ctc_best_path(logits=logits, vocab=self.vocab, blank=len(self.vocab))
+
+
+class CRNN(Engine):
+    """CRNN Onnx loader
+
+    Args:
+    ----
+        model_path: path or url to onnx model file
+        vocab: vocabulary used for encoding
+        cfg: configuration dictionary
+        **kwargs: additional arguments to be passed to `Engine`
+    """
+
+    _children_names: List[str] = ["postprocessor"]
+
+    def __init__(
+        self,
+        model_path: str,
+        vocab: str,
+        cfg: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(url=model_path, **kwargs)
+        self.vocab = vocab
+        self.cfg = cfg
+        self.postprocessor = CRNNPostProcessor(self.vocab)
+
+    def __call__(
+        self,
+        x: np.ndarray,
+        return_model_output: bool = False,
+    ) -> Dict[str, Any]:
+        logits = self.run(x)
+
+        out: Dict[str, Any] = {}
+        if return_model_output:
+            out["out_map"] = logits
+
+        # Post-process
+        out["preds"] = self.postprocessor(logits)
+
+        return out
+
+
+def _crnn(
+    arch: str,
+    model_path: str,
+    **kwargs: Any,
+) -> CRNN:
+    kwargs["vocab"] = kwargs.get("vocab", default_cfgs[arch]["vocab"])
+
+    _cfg = deepcopy(default_cfgs[arch])
+    _cfg["vocab"] = kwargs["vocab"]
+    _cfg["input_shape"] = kwargs.get("input_shape", default_cfgs[arch]["input_shape"])
+
+    # Build the model
+    return CRNN(model_path, cfg=_cfg, **kwargs)
+
+
+def crnn_vgg16_bn(model_path: str = default_cfgs["crnn_vgg16_bn"]["url"], **kwargs: Any) -> CRNN:
+    """CRNN with a VGG-16 backbone as described in `"An End-to-End Trainable Neural Network for Image-based
+    Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import crnn_vgg16_bn
+    >>> model = crnn_vgg16_bn()
+    >>> input_tensor = np.random.rand(1, 3, 32, 128)
+    >>> out = model(input_tensor)
+
+    Args:
+    ----
+        model_path: path to onnx model file, defaults to url in default_cfgs
+        **kwargs: keyword arguments of the CRNN architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _crnn("crnn_vgg16_bn", model_path, **kwargs)
+
+
+def crnn_mobilenet_v3_small(model_path: str = default_cfgs["crnn_mobilenet_v3_small"]["url"], **kwargs: Any) -> CRNN:
+    """CRNN with a MobileNet V3 Small backbone as described in `"An End-to-End Trainable Neural Network for Image-based
+    Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import crnn_mobilenet_v3_small
+    >>> model = crnn_mobilenet_v3_small()
+    >>> input_tensor = np.random.rand(1, 3, 32, 128)
+    >>> out = model(input_tensor)
+
+    Args:
+    ----
+        model_path: path to onnx model file, defaults to url in default_cfgs
+        **kwargs: keyword arguments of the CRNN architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _crnn("crnn_mobilenet_v3_small", model_path, **kwargs)
+
+
+def crnn_mobilenet_v3_large(model_path: str = default_cfgs["crnn_mobilenet_v3_large"]["url"], **kwargs: Any) -> CRNN:
+    """CRNN with a MobileNet V3 Large backbone as described in `"An End-to-End Trainable Neural Network for Image-based
+    Sequence Recognition and Its Application to Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import crnn_mobilenet_v3_large
+    >>> model = crnn_mobilenet_v3_large()
+    >>> input_tensor = np.random.rand(1, 3, 32, 128)
+    >>> out = model(input_tensor)
+
+    Args:
+    ----
+        model_path: path to onnx model file, defaults to url in default_cfgs
+        **kwargs: keyword arguments of the CRNN architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _crnn("crnn_mobilenet_v3_large", model_path, **kwargs)

onnxtr/models/recognition/models/master.py

@@ -0,0 +1,145 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from copy import deepcopy
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+from scipy.special import softmax
+
+from onnxtr.utils import VOCABS
+
+from ...engine import Engine
+from ..core import RecognitionPostProcessor
+
+__all__ = ["MASTER", "master"]
+
+
+default_cfgs: Dict[str, Dict[str, Any]] = {
+    "master": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 32, 128),
+        "vocab": VOCABS["french"],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/master-b1287fcd.onnx",
+    },
+}
+
+
+class MASTER(Engine):
+    """MASTER Onnx loader
+
+    Args:
+    ----
+        model_path: path or url to onnx model file
+        vocab: vocabulary, (without EOS, SOS, PAD)
+        cfg: dictionary containing information about the model
+        **kwargs: additional arguments to be passed to `Engine`
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        vocab: str,
+        cfg: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(url=model_path, **kwargs)
+
+        self.vocab = vocab
+        self.cfg = cfg
+        self.postprocessor = MASTERPostProcessor(vocab=self.vocab)
+
+    def __call__(
+        self,
+        x: np.ndarray,
+        return_model_output: bool = False,
+    ) -> Dict[str, Any]:
+        """Call function
+
+        Args:
+        ----
+            x: images
+            return_model_output: if True, return logits
+
+        Returns:
+        -------
+            A dictionnary containing eventually logits and predictions.
+        """
+        logits = self.run(x)
+        out: Dict[str, Any] = {}
+
+        if return_model_output:
+            out["out_map"] = logits
+
+        out["preds"] = self.postprocessor(logits)
+
+        return out
+
+
+class MASTERPostProcessor(RecognitionPostProcessor):
+    """Post-processor for the MASTER model
+
+    Args:
+    ----
+        vocab: string containing the ordered sequence of supported characters
+    """
+
+    def __init__(
+        self,
+        vocab: str,
+    ) -> None:
+        super().__init__(vocab)
+        self._embedding = list(vocab) + ["<eos>"] + ["<sos>"] + ["<pad>"]
+
+    def __call__(self, logits: np.ndarray) -> List[Tuple[str, float]]:
+        # compute pred with argmax for attention models
+        out_idxs = np.argmax(logits, axis=-1)
+        # N x L
+        probs = np.take_along_axis(softmax(logits, axis=-1), out_idxs[..., None], axis=-1).squeeze(-1)
+        # Take the minimum confidence of the sequence
+        probs = np.min(probs, axis=1)
+
+        word_values = [
+            "".join(self._embedding[idx] for idx in encoded_seq).split("<eos>")[0] for encoded_seq in out_idxs
+        ]
+
+        return list(zip(word_values, np.clip(probs, 0, 1).astype(float).tolist()))
+
+
+def _master(
+    arch: str,
+    model_path: str,
+    **kwargs: Any,
+) -> MASTER:
+    # Patch the config
+    _cfg = deepcopy(default_cfgs[arch])
+    _cfg["input_shape"] = kwargs.get("input_shape", _cfg["input_shape"])
+    _cfg["vocab"] = kwargs.get("vocab", _cfg["vocab"])
+
+    kwargs["vocab"] = _cfg["vocab"]
+
+    return MASTER(model_path, cfg=_cfg, **kwargs)
+
+
+def master(model_path: str = default_cfgs["master"]["url"], **kwargs: Any) -> MASTER:
+    """MASTER as described in paper: <https://arxiv.org/pdf/1910.02562.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import master
+    >>> model = master()
+    >>> input_tensor = np.random.rand(1, 3, 32, 128)
+    >>> out = model(input_tensor)
+
+    Args:
+    ----
+        model_path: path to onnx model file, defaults to url in default_cfgs
+        **kwargs: keywoard arguments passed to the MASTER architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _master("master", model_path, **kwargs)

onnxtr/models/recognition/models/parseq.py

@@ -0,0 +1,134 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from copy import deepcopy
+from typing import Any, Dict, Optional
+
+import numpy as np
+from scipy.special import softmax
+
+from onnxtr.utils import VOCABS
+
+from ...engine import Engine
+from ..core import RecognitionPostProcessor
+
+__all__ = ["PARSeq", "parseq"]
+
+default_cfgs: Dict[str, Dict[str, Any]] = {
+    "parseq": {
+        "mean": (0.694, 0.695, 0.693),
+        "std": (0.299, 0.296, 0.301),
+        "input_shape": (3, 32, 128),
+        "vocab": VOCABS["french"],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.0.1/parseq-00b40714.onnx",
+    },
+}
+
+
+class PARSeq(Engine):
+    """PARSeq Onnx loader
+
+    Args:
+    ----
+        vocab: vocabulary used for encoding
+        cfg: dictionary containing information about the model
+        **kwargs: additional arguments to be passed to `Engine`
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        vocab: str,
+        cfg: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(url=model_path, **kwargs)
+        self.vocab = vocab
+        self.cfg = cfg
+        self.postprocessor = PARSeqPostProcessor(vocab=self.vocab)
+
+    def __call__(
+        self,
+        x: np.ndarray,
+        return_model_output: bool = False,
+    ) -> Dict[str, Any]:
+        logits = self.run(x)
+        out: Dict[str, Any] = {}
+
+        if return_model_output:
+            out["out_map"] = logits
+
+        out["preds"] = self.postprocessor(logits)
+        return out
+
+
+class PARSeqPostProcessor(RecognitionPostProcessor):
+    """Post processor for PARSeq architecture
+
+    Args:
+    ----
+        vocab: string containing the ordered sequence of supported characters
+    """
+
+    def __init__(
+        self,
+        vocab: str,
+    ) -> None:
+        super().__init__(vocab)
+        self._embedding = list(vocab) + ["<eos>", "<sos>", "<pad>"]
+
+    def __call__(self, logits):
+        # compute pred with argmax for attention models
+        out_idxs = np.argmax(logits, axis=-1)
+        preds_prob = softmax(logits, axis=-1).max(axis=-1)
+
+        word_values = [
+            "".join(self._embedding[idx] for idx in encoded_seq).split("<eos>")[0] for encoded_seq in out_idxs
+        ]
+        # compute probabilties for each word up to the EOS token
+        probs = [
+            preds_prob[i, : len(word)].clip(0, 1).mean().astype(float) if word else 0.0
+            for i, word in enumerate(word_values)
+        ]
+
+        return list(zip(word_values, probs))
+
+
+def _parseq(
+    arch: str,
+    model_path: str,
+    **kwargs: Any,
+) -> PARSeq:
+    # Patch the config
+    _cfg = deepcopy(default_cfgs[arch])
+    _cfg["vocab"] = kwargs.get("vocab", _cfg["vocab"])
+    _cfg["input_shape"] = kwargs.get("input_shape", _cfg["input_shape"])
+
+    kwargs["vocab"] = _cfg["vocab"]
+
+    # Build the model
+    return PARSeq(model_path, cfg=_cfg, **kwargs)
+
+
+def parseq(model_path: str = default_cfgs["parseq"]["url"], **kwargs: Any) -> PARSeq:
+    """PARSeq architecture from
+    `"Scene Text Recognition with Permuted Autoregressive Sequence Models" <https://arxiv.org/pdf/2207.06966>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import parseq
+    >>> model = parseq()
+    >>> input_tensor = np.random.rand(1, 3, 32, 128)
+    >>> out = model(input_tensor)
+
+    Args:
+    ----
+        model_path: path to onnx model file, defaults to url in default_cfgs
+        **kwargs: keyword arguments of the PARSeq architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _parseq("parseq", model_path, **kwargs)