onnxtr 0.1.0__py3-none-any.whl

onnxtr/models/recognition/models/sar.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__ = ["SAR", "sar_resnet31"]
+
+default_cfgs: Dict[str, Dict[str, Any]] = {
+    "sar_resnet31": {
+        "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/sar_resnet31-395f8005.onnx",
+    },
+}
+
+
+class SAR(Engine):
+    """SAR Onnx loader
+
+    Args:
+    ----
+        model_path: path to onnx model file
+        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 = SARPostProcessor(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 SARPostProcessor(RecognitionPostProcessor):
+    """Post processor for SAR architectures
+
+    Args:
+    ----
+        embedding: string containing the ordered sequence of supported characters
+    """
+
+    def __init__(
+        self,
+        vocab: str,
+    ) -> None:
+        super().__init__(vocab)
+        self._embedding = list(self.vocab) + ["<eos>"]
+
+    def __call__(self, logits):
+        # 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 _sar(
+    arch: str,
+    model_path: str,
+    **kwargs: Any,
+) -> SAR:
+    # 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 SAR(model_path, cfg=_cfg, **kwargs)
+
+
+def sar_resnet31(model_path: str = default_cfgs["sar_resnet31"]["url"], **kwargs: Any) -> SAR:
+    """SAR with a resnet-31 feature extractor as described in `"Show, Attend and Read:A Simple and Strong
+    Baseline for Irregular Text Recognition" <https://arxiv.org/pdf/1811.00751.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import sar_resnet31
+    >>> model = sar_resnet31()
+    >>> 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 SAR architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _sar("sar_resnet31", model_path, **kwargs)

onnxtr/models/recognition/models/vitstr.py

@@ -0,0 +1,166 @@
+# 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__ = ["ViTSTR", "vitstr_small", "vitstr_base"]
+
+default_cfgs: Dict[str, Dict[str, Any]] = {
+    "vitstr_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/vitstr_small-3ff9c500.onnx",
+    },
+    "vitstr_base": {
+        "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/vitstr_base-ff62f5be.onnx",
+    },
+}
+
+
+class ViTSTR(Engine):
+    """ViTSTR Onnx loader
+
+    Args:
+    ----
+        model_path: path to onnx model file
+        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 = ViTSTRPostProcessor(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 ViTSTRPostProcessor(RecognitionPostProcessor):
+    """Post processor for ViTSTR 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>"]
+
+    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 _vitstr(
+    arch: str,
+    model_path: str,
+    **kwargs: Any,
+) -> ViTSTR:
+    # 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 ViTSTR(model_path, cfg=_cfg, **kwargs)
+
+
+def vitstr_small(model_path: str = default_cfgs["vitstr_small"]["url"], **kwargs: Any) -> ViTSTR:
+    """ViTSTR-Small as described in `"Vision Transformer for Fast and Efficient Scene Text Recognition"
+    <https://arxiv.org/pdf/2105.08582.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import vitstr_small
+    >>> model = vitstr_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 ViTSTR architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _vitstr("vitstr_small", model_path, **kwargs)
+
+
+def vitstr_base(model_path: str = default_cfgs["vitstr_base"]["url"], **kwargs: Any) -> ViTSTR:
+    """ViTSTR-Base as described in `"Vision Transformer for Fast and Efficient Scene Text Recognition"
+    <https://arxiv.org/pdf/2105.08582.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import vitstr_base
+    >>> model = vitstr_base()
+    >>> 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 ViTSTR architecture
+
+    Returns:
+    -------
+        text recognition architecture
+    """
+    return _vitstr("vitstr_base", model_path, **kwargs)

onnxtr/models/recognition/predictor/__init__.py

@@ -0,0 +1 @@
+from .base import *

onnxtr/models/recognition/predictor/_utils.py

@@ -0,0 +1,86 @@
+# 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 typing import List, Tuple, Union
+
+import numpy as np
+
+from ..utils import merge_multi_strings
+
+__all__ = ["split_crops", "remap_preds"]
+
+
+def split_crops(
+    crops: List[np.ndarray],
+    max_ratio: float,
+    target_ratio: int,
+    dilation: float,
+    channels_last: bool = True,
+) -> Tuple[List[np.ndarray], List[Union[int, Tuple[int, int]]], bool]:
+    """Chunk crops horizontally to match a given aspect ratio
+
+    Args:
+    ----
+        crops: list of numpy array of shape (H, W, 3) if channels_last or (3, H, W) otherwise
+        max_ratio: the maximum aspect ratio that won't trigger the chunk
+        target_ratio: when crops are chunked, they will be chunked to match this aspect ratio
+        dilation: the width dilation of final chunks (to provide some overlaps)
+        channels_last: whether the numpy array has dimensions in channels last order
+
+    Returns:
+    -------
+        a tuple with the new crops, their mapping, and a boolean specifying whether any remap is required
+    """
+    _remap_required = False
+    crop_map: List[Union[int, Tuple[int, int]]] = []
+    new_crops: List[np.ndarray] = []
+    for crop in crops:
+        h, w = crop.shape[:2] if channels_last else crop.shape[-2:]
+        aspect_ratio = w / h
+        if aspect_ratio > max_ratio:
+            # Determine the number of crops, reference aspect ratio = 4 = 128 / 32
+            num_subcrops = int(aspect_ratio // target_ratio)
+            # Find the new widths, additional dilation factor to overlap crops
+            width = dilation * w / num_subcrops
+            centers = [(w / num_subcrops) * (1 / 2 + idx) for idx in range(num_subcrops)]
+            # Get the crops
+            if channels_last:
+                _crops = [
+                    crop[:, max(0, int(round(center - width / 2))) : min(w - 1, int(round(center + width / 2))), :]
+                    for center in centers
+                ]
+            else:
+                _crops = [
+                    crop[:, :, max(0, int(round(center - width / 2))) : min(w - 1, int(round(center + width / 2)))]
+                    for center in centers
+                ]
+            # Avoid sending zero-sized crops
+            _crops = [crop for crop in _crops if all(s > 0 for s in crop.shape)]
+            # Record the slice of crops
+            crop_map.append((len(new_crops), len(new_crops) + len(_crops)))
+            new_crops.extend(_crops)
+            # At least one crop will require merging
+            _remap_required = True
+        else:
+            crop_map.append(len(new_crops))
+            new_crops.append(crop)
+
+    return new_crops, crop_map, _remap_required
+
+
+def remap_preds(
+    preds: List[Tuple[str, float]], crop_map: List[Union[int, Tuple[int, int]]], dilation: float
+) -> List[Tuple[str, float]]:
+    remapped_out = []
+    for _idx in crop_map:
+        # Crop hasn't been split
+        if isinstance(_idx, int):
+            remapped_out.append(preds[_idx])
+        else:
+            # unzip
+            vals, probs = zip(*preds[_idx[0] : _idx[1]])
+            # Merge the string values
+            remapped_out.append((merge_multi_strings(vals, dilation), min(probs)))  # type: ignore[arg-type]
+    return remapped_out

onnxtr/models/recognition/predictor/base.py

@@ -0,0 +1,79 @@
+# 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 typing import Any, List, Sequence, Tuple
+
+import numpy as np
+
+from onnxtr.models.preprocessor import PreProcessor
+from onnxtr.utils.repr import NestedObject
+
+from ._utils import remap_preds, split_crops
+
+__all__ = ["RecognitionPredictor"]
+
+
+class RecognitionPredictor(NestedObject):
+    """Implements an object able to identify character sequences in images
+
+    Args:
+    ----
+        pre_processor: transform inputs for easier batched model inference
+        model: core recognition architecture
+        split_wide_crops: wether to use crop splitting for high aspect ratio crops
+    """
+
+    def __init__(
+        self,
+        pre_processor: PreProcessor,
+        model: Any,
+        split_wide_crops: bool = True,
+    ) -> None:
+        super().__init__()
+        self.pre_processor = pre_processor
+        self.model = model
+        self.split_wide_crops = split_wide_crops
+        self.critical_ar = 8  # Critical aspect ratio
+        self.dil_factor = 1.4  # Dilation factor to overlap the crops
+        self.target_ar = 6  # Target aspect ratio
+
+    def __call__(
+        self,
+        crops: Sequence[np.ndarray],
+        **kwargs: Any,
+    ) -> List[Tuple[str, float]]:
+        if len(crops) == 0:
+            return []
+        # Dimension check
+        if any(crop.ndim != 3 for crop in crops):
+            raise ValueError("incorrect input shape: all crops are expected to be multi-channel 2D images.")
+
+        # Split crops that are too wide
+        remapped = False
+        if self.split_wide_crops:
+            new_crops, crop_map, remapped = split_crops(
+                crops,  # type: ignore[arg-type]
+                self.critical_ar,
+                self.target_ar,
+                self.dil_factor,
+                True,
+            )
+            if remapped:
+                crops = new_crops
+
+        # Resize & batch them
+        processed_batches = self.pre_processor(crops)  # type: ignore[arg-type]
+
+        # Forward it
+        raw = [self.model(batch, **kwargs)["preds"] for batch in processed_batches]
+
+        # Process outputs
+        out = [charseq for batch in raw for charseq in batch]
+
+        # Remap crops
+        if self.split_wide_crops and remapped:
+            out = remap_preds(out, crop_map, self.dil_factor)
+
+        return out

onnxtr/models/recognition/utils.py

@@ -0,0 +1,89 @@
+# 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 typing import List
+
+from rapidfuzz.distance import Levenshtein
+
+__all__ = ["merge_strings", "merge_multi_strings"]
+
+
+def merge_strings(a: str, b: str, dil_factor: float) -> str:
+    """Merges 2 character sequences in the best way to maximize the alignment of their overlapping characters.
+
+    Args:
+    ----
+        a: first char seq, suffix should be similar to b's prefix.
+        b: second char seq, prefix should be similar to a's suffix.
+        dil_factor: dilation factor of the boxes to overlap, should be > 1. This parameter is
+            only used when the mother sequence is splitted on a character repetition
+
+    Returns:
+    -------
+        A merged character sequence.
+
+    Example::
+        >>> from onnxtr.model.recognition.utils import merge_sequences
+        >>> merge_sequences('abcd', 'cdefgh', 1.4)
+        'abcdefgh'
+        >>> merge_sequences('abcdi', 'cdefgh', 1.4)
+        'abcdefgh'
+    """
+    seq_len = min(len(a), len(b))
+    if seq_len == 0:  # One sequence is empty, return the other
+        return b if len(a) == 0 else a
+
+    # Initialize merging index and corresponding score (mean Levenstein)
+    min_score, index = 1.0, 0  # No overlap, just concatenate
+
+    scores = [Levenshtein.distance(a[-i:], b[:i], processor=None) / i for i in range(1, seq_len + 1)]
+
+    # Edge case (split in the middle of char repetitions): if it starts with 2 or more 0
+    if len(scores) > 1 and (scores[0], scores[1]) == (0, 0):
+        # Compute n_overlap (number of overlapping chars, geometrically determined)
+        n_overlap = round(len(b) * (dil_factor - 1) / dil_factor)
+        # Find the number of consecutive zeros in the scores list
+        # Impossible to have a zero after a non-zero score in that case
+        n_zeros = sum(val == 0 for val in scores)
+        # Index is bounded by the geometrical overlap to avoid collapsing repetitions
+        min_score, index = 0, min(n_zeros, n_overlap)
+
+    else:  # Common case: choose the min score index
+        for i, score in enumerate(scores):
+            if score < min_score:
+                min_score, index = score, i + 1  # Add one because first index is an overlap of 1 char
+
+    # Merge with correct overlap
+    if index == 0:
+        return a + b
+    return a[:-1] + b[index - 1 :]
+
+
+def merge_multi_strings(seq_list: List[str], dil_factor: float) -> str:
+    """Recursively merges consecutive string sequences with overlapping characters.
+
+    Args:
+    ----
+        seq_list: list of sequences to merge. Sequences need to be ordered from left to right.
+        dil_factor: dilation factor of the boxes to overlap, should be > 1. This parameter is
+            only used when the mother sequence is splitted on a character repetition
+
+    Returns:
+    -------
+        A merged character sequence
+
+    Example::
+        >>> from onnxtr.model.recognition.utils import merge_multi_sequences
+        >>> merge_multi_sequences(['abc', 'bcdef', 'difghi', 'aijkl'], 1.4)
+        'abcdefghijkl'
+    """
+
+    def _recursive_merge(a: str, seq_list: List[str], dil_factor: float) -> str:
+        # Recursive version of compute_overlap
+        if len(seq_list) == 1:
+            return merge_strings(a, seq_list[0], dil_factor)
+        return _recursive_merge(merge_strings(a, seq_list[0], dil_factor), seq_list[1:], dil_factor)
+
+    return _recursive_merge("", seq_list, dil_factor)

onnxtr/models/recognition/zoo.py

@@ -0,0 +1,69 @@
+# 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 typing import Any, List
+
+from onnxtr.models.preprocessor import PreProcessor
+
+from .. import recognition
+from .predictor import RecognitionPredictor
+
+__all__ = ["recognition_predictor"]
+
+
+ARCHS: List[str] = [
+    "crnn_vgg16_bn",
+    "crnn_mobilenet_v3_small",
+    "crnn_mobilenet_v3_large",
+    "sar_resnet31",
+    "master",
+    "vitstr_small",
+    "vitstr_base",
+    "parseq",
+]
+
+
+def _predictor(arch: Any, **kwargs: Any) -> RecognitionPredictor:
+    if isinstance(arch, str):
+        if arch not in ARCHS:
+            raise ValueError(f"unknown architecture '{arch}'")
+
+        _model = recognition.__dict__[arch]()
+    else:
+        if not isinstance(
+            arch, (recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq)
+        ):
+            raise ValueError(f"unknown architecture: {type(arch)}")
+        _model = arch
+
+    kwargs["mean"] = kwargs.get("mean", _model.cfg["mean"])
+    kwargs["std"] = kwargs.get("std", _model.cfg["std"])
+    kwargs["batch_size"] = kwargs.get("batch_size", 1024)
+    input_shape = _model.cfg["input_shape"][1:]
+    predictor = RecognitionPredictor(PreProcessor(input_shape, preserve_aspect_ratio=True, **kwargs), _model)
+
+    return predictor
+
+
+def recognition_predictor(arch: Any = "crnn_vgg16_bn", **kwargs: Any) -> RecognitionPredictor:
+    """Text recognition architecture.
+
+    Example::
+        >>> import numpy as np
+        >>> from onnxtr.models import recognition_predictor
+        >>> model = recognition_predictor()
+        >>> input_page = (255 * np.random.rand(32, 128, 3)).astype(np.uint8)
+        >>> out = model([input_page])
+
+    Args:
+    ----
+        arch: name of the architecture or model itself to use (e.g. 'crnn_vgg16_bn')
+        **kwargs: optional parameters to be passed to the architecture
+
+    Returns:
+    -------
+        Recognition predictor
+    """
+    return _predictor(arch, **kwargs)