onnxtr 0.6.3__tar.gz → 0.7.1__tar.gz

{onnxtr-0.6.3 → onnxtr-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxtr
-Version: 0.6.3
+Version: 0.7.1
 Summary: Onnx Text Recognition (OnnxTR): docTR Onnx-Wrapper for high-performance OCR on documents.
 Author-email: Felix Dittrich <felixdittrich92@gmail.com>
 Maintainer: Felix Dittrich
@@ -292,8 +292,8 @@ Dynamic: license-file
 [![codecov](https://codecov.io/gh/felixdittrich92/OnnxTR/graph/badge.svg?token=WVFRCQBOLI)](https://codecov.io/gh/felixdittrich92/OnnxTR)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fff4d764bb14fb8b4f4afeb9587231b)](https://app.codacy.com/gh/felixdittrich92/OnnxTR/dashboard?utm_source=gh&utm_medium=referral&utm_content=&utm_campaign=Badge_grade)
 [![CodeFactor](https://www.codefactor.io/repository/github/felixdittrich92/onnxtr/badge)](https://www.codefactor.io/repository/github/felixdittrich92/onnxtr)
-[![Socket Badge](https://socket.dev/api/badge/pypi/package/onnxtr/0.6.2?artifact_id=tar-gz)](https://socket.dev/pypi/package/onnxtr/overview/0.6.2/tar-gz)
-[![Pypi](https://img.shields.io/badge/pypi-v0.6.3-blue.svg)](https://pypi.org/project/OnnxTR/)
+[![Socket Badge](https://socket.dev/api/badge/pypi/package/onnxtr/0.7.1?artifact_id=tar-gz)](https://socket.dev/pypi/package/onnxtr/overview/0.7.1/tar-gz)
+[![Pypi](https://img.shields.io/badge/pypi-v0.7.1-blue.svg)](https://pypi.org/project/OnnxTR/)
 [![Docker Images](https://img.shields.io/badge/Docker-4287f5?style=flat&logo=docker&logoColor=white)](https://github.com/felixdittrich92/OnnxTR/pkgs/container/onnxtr)
 [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Felix92/OnnxTR-OCR)
 ![PyPI - Downloads](https://img.shields.io/pypi/dm/onnxtr)
@@ -575,6 +575,7 @@ Credits where it's due: this repository provides ONNX models for the following a
 - MASTER: [MASTER: Multi-Aspect Non-local Network for Scene Text Recognition](https://arxiv.org/pdf/1910.02562.pdf).
 - ViTSTR: [Vision Transformer for Fast and Efficient Scene Text Recognition](https://arxiv.org/pdf/2105.08582.pdf).
 - PARSeq: [Scene Text Recognition with Permuted Autoregressive Sequence Models](https://arxiv.org/pdf/2207.06966).
+- VIPTR: [A Vision Permutable Extractor for Fast and Efficient Scene Text Recognition](https://arxiv.org/abs/2401.10110).
 
 ```python
 predictor = ocr_predictor()
@@ -602,6 +603,7 @@ predictor.list_archs()
             'vitstr_small',
             'vitstr_base',
             'parseq'
+            'viptr_tiny',  # No 8-bit support
         ]
 }
 ```

{onnxtr-0.6.3 → onnxtr-0.7.1}/README.md

@@ -7,8 +7,8 @@
 [![codecov](https://codecov.io/gh/felixdittrich92/OnnxTR/graph/badge.svg?token=WVFRCQBOLI)](https://codecov.io/gh/felixdittrich92/OnnxTR)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fff4d764bb14fb8b4f4afeb9587231b)](https://app.codacy.com/gh/felixdittrich92/OnnxTR/dashboard?utm_source=gh&utm_medium=referral&utm_content=&utm_campaign=Badge_grade)
 [![CodeFactor](https://www.codefactor.io/repository/github/felixdittrich92/onnxtr/badge)](https://www.codefactor.io/repository/github/felixdittrich92/onnxtr)
-[![Socket Badge](https://socket.dev/api/badge/pypi/package/onnxtr/0.6.2?artifact_id=tar-gz)](https://socket.dev/pypi/package/onnxtr/overview/0.6.2/tar-gz)
-[![Pypi](https://img.shields.io/badge/pypi-v0.6.3-blue.svg)](https://pypi.org/project/OnnxTR/)
+[![Socket Badge](https://socket.dev/api/badge/pypi/package/onnxtr/0.7.1?artifact_id=tar-gz)](https://socket.dev/pypi/package/onnxtr/overview/0.7.1/tar-gz)
+[![Pypi](https://img.shields.io/badge/pypi-v0.7.1-blue.svg)](https://pypi.org/project/OnnxTR/)
 [![Docker Images](https://img.shields.io/badge/Docker-4287f5?style=flat&logo=docker&logoColor=white)](https://github.com/felixdittrich92/OnnxTR/pkgs/container/onnxtr)
 [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Felix92/OnnxTR-OCR)
 ![PyPI - Downloads](https://img.shields.io/pypi/dm/onnxtr)
@@ -290,6 +290,7 @@ Credits where it's due: this repository provides ONNX models for the following a
 - MASTER: [MASTER: Multi-Aspect Non-local Network for Scene Text Recognition](https://arxiv.org/pdf/1910.02562.pdf).
 - ViTSTR: [Vision Transformer for Fast and Efficient Scene Text Recognition](https://arxiv.org/pdf/2105.08582.pdf).
 - PARSeq: [Scene Text Recognition with Permuted Autoregressive Sequence Models](https://arxiv.org/pdf/2207.06966).
+- VIPTR: [A Vision Permutable Extractor for Fast and Efficient Scene Text Recognition](https://arxiv.org/abs/2401.10110).
 
 ```python
 predictor = ocr_predictor()
@@ -317,6 +318,7 @@ predictor.list_archs()
             'vitstr_small',
             'vitstr_base',
             'parseq'
+            'viptr_tiny',  # No 8-bit support
         ]
 }
 ```

{onnxtr-0.6.3 → onnxtr-0.7.1}/onnxtr/models/predictor/predictor.py

@@ -115,7 +115,7 @@ class OCRPredictor(NestedObject, _OCRPredictor):
         # Crop images
         crops, loc_preds = self._prepare_crops(
             pages,
-            loc_preds,  # type: ignore[arg-type]
+            loc_preds,
             channels_last=True,
             assume_straight_pages=self.assume_straight_pages,
             assume_horizontal=self._page_orientation_disabled,

{onnxtr-0.6.3 → onnxtr-0.7.1}/onnxtr/models/recognition/models/__init__.py

@@ -3,3 +3,4 @@ from .sar import *
 from .master import *
 from .vitstr import *
 from .parseq import *
+from .viptr import *

onnxtr-0.7.1/onnxtr/models/recognition/models/viptr.py

@@ -0,0 +1,179 @@
+# Copyright (C) 2021-2025, 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 logging
+from copy import deepcopy
+from itertools import groupby
+from typing import Any
+
+import numpy as np
+from scipy.special import softmax
+
+from onnxtr.utils import VOCABS
+
+from ...engine import Engine, EngineConfig
+from ..core import RecognitionPostProcessor
+
+__all__ = ["VIPTR", "viptr_tiny"]
+
+default_cfgs: dict[str, dict[str, Any]] = {
+    "viptr_tiny": {
+        "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.6.3/viptr_tiny-499b8015.onnx",
+        "url_8_bit": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.6.3/viptr_tiny-499b8015.onnx",
+    },
+}
+
+
+class VIPTRPostProcessor(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 VIPTR(Engine):
+    """VIPTR Onnx loader
+
+    Args:
+        model_path: path or url to onnx model file
+        vocab: vocabulary used for encoding
+        engine_cfg: configuration for the inference engine
+        cfg: configuration dictionary
+        **kwargs: additional arguments to be passed to `Engine`
+    """
+
+    _children_names: list[str] = ["postprocessor"]
+
+    def __init__(
+        self,
+        model_path: str,
+        vocab: str,
+        engine_cfg: EngineConfig | None = None,
+        cfg: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
+        self.vocab = vocab
+        self.cfg = cfg
+
+        self.postprocessor = VIPTRPostProcessor(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 _viptr(
+    arch: str,
+    model_path: str,
+    load_in_8_bit: bool = False,
+    engine_cfg: EngineConfig | None = None,
+    **kwargs: Any,
+) -> VIPTR:
+    if load_in_8_bit:
+        logging.warning("VIPTR models do not support 8-bit quantization yet. Loading full precision model...")
+    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"])
+    # Patch the url
+    model_path = default_cfgs[arch]["url_8_bit"] if load_in_8_bit and "http" in model_path else model_path
+
+    # Build the model
+    return VIPTR(model_path, cfg=_cfg, engine_cfg=engine_cfg, **kwargs)
+
+
+def viptr_tiny(
+    model_path: str = default_cfgs["viptr_tiny"]["url"],
+    load_in_8_bit: bool = False,
+    engine_cfg: EngineConfig | None = None,
+    **kwargs: Any,
+) -> VIPTR:
+    """VIPTR as described in `"A Vision Permutable Extractor for Fast and Efficient
+    Scene Text Recognition" <https://arxiv.org/pdf/1507.05717.pdf>`_.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import viptr_tiny
+    >>> model = viptr_tiny()
+    >>> 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
+        load_in_8_bit: whether to load the the 8-bit quantized model, defaults to False
+        engine_cfg: configuration for the inference engine
+        **kwargs: keyword arguments of the VIPTR architecture
+
+    Returns:
+        text recognition architecture
+    """
+    return _viptr("viptr_tiny", model_path, load_in_8_bit, engine_cfg, **kwargs)

onnxtr-0.7.1/onnxtr/models/recognition/predictor/_utils.py

@@ -0,0 +1,145 @@
+# Copyright (C) 2021-2025, 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
+
+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,
+    split_overlap_ratio: float,
+    channels_last: bool = True,
+) -> tuple[list[np.ndarray], list[int | tuple[int, int, float]], bool]:
+    """
+    Split crops horizontally if they exceed a given aspect ratio.
+
+    Args:
+        crops: List of image crops (H, W, C) if channels_last else (C, H, W).
+        max_ratio: Aspect ratio threshold above which crops are split.
+        target_ratio: Target aspect ratio after splitting (e.g., 4 for 128x32).
+        split_overlap_ratio: Desired overlap between splits (as a fraction of split width).
+        channels_last: Whether the crops are in channels-last format.
+
+    Returns:
+        A tuple containing:
+            - The new list of crops (possibly with splits),
+            - A mapping indicating how to reassemble predictions,
+            - A boolean indicating whether remapping is required.
+    """
+    if split_overlap_ratio <= 0.0 or split_overlap_ratio >= 1.0:
+        raise ValueError(f"Valid range for split_overlap_ratio is (0.0, 1.0), but is: {split_overlap_ratio}")
+
+    remap_required = False
+    new_crops: list[np.ndarray] = []
+    crop_map: list[int | tuple[int, int, float]] = []
+
+    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:
+            split_width = max(1, math.ceil(h * target_ratio))
+            overlap_width = max(0, math.floor(split_width * split_overlap_ratio))
+
+            splits, last_overlap = _split_horizontally(crop, split_width, overlap_width, channels_last)
+
+            # Remove any empty splits
+            splits = [s for s in splits if all(dim > 0 for dim in s.shape)]
+            if splits:
+                crop_map.append((len(new_crops), len(new_crops) + len(splits), last_overlap))
+                new_crops.extend(splits)
+                remap_required = True
+            else:
+                # Fallback: treat it as a single crop
+                crop_map.append(len(new_crops))
+                new_crops.append(crop)
+        else:
+            crop_map.append(len(new_crops))
+            new_crops.append(crop)
+
+    return new_crops, crop_map, remap_required
+
+
+def _split_horizontally(
+    image: np.ndarray, split_width: int, overlap_width: int, channels_last: bool
+) -> tuple[list[np.ndarray], float]:
+    """
+    Horizontally split a single image with overlapping regions.
+
+    Args:
+        image: The image to split (H, W, C) if channels_last else (C, H, W).
+        split_width: Width of each split.
+        overlap_width: Width of the overlapping region.
+        channels_last: Whether the image is in channels-last format.
+
+    Returns:
+        - A list of horizontal image slices.
+        - The actual overlap ratio of the last split.
+    """
+    image_width = image.shape[1] if channels_last else image.shape[-1]
+    if image_width <= split_width:
+        return [image], 0.0
+
+    # Compute start columns for each split
+    step = split_width - overlap_width
+    starts = list(range(0, image_width - split_width + 1, step))
+
+    # Ensure the last patch reaches the end of the image
+    if starts[-1] + split_width < image_width:
+        starts.append(image_width - split_width)
+
+    splits = []
+    for start_col in starts:
+        end_col = start_col + split_width
+        if channels_last:
+            split = image[:, start_col:end_col, :]
+        else:
+            split = image[:, :, start_col:end_col]
+        splits.append(split)
+
+    # Calculate the last overlap ratio, if only one split no overlap
+    last_overlap = 0
+    if len(starts) > 1:
+        last_overlap = (starts[-2] + split_width) - starts[-1]
+    last_overlap_ratio = last_overlap / split_width if split_width else 0.0
+
+    return splits, last_overlap_ratio
+
+
+def remap_preds(
+    preds: list[tuple[str, float]],
+    crop_map: list[int | tuple[int, int, float]],
+    overlap_ratio: float,
+) -> list[tuple[str, float]]:
+    """
+    Reconstruct predictions from possibly split crops.
+
+    Args:
+        preds: List of (text, confidence) tuples from each crop.
+        crop_map: Map returned by `split_crops`.
+        overlap_ratio: Overlap ratio used during splitting.
+
+    Returns:
+        List of merged (text, confidence) tuples corresponding to original crops.
+    """
+    remapped = []
+    for item in crop_map:
+        if isinstance(item, int):
+            remapped.append(preds[item])
+        else:
+            start_idx, end_idx, last_overlap = item
+            text_parts, confidences = zip(*preds[start_idx:end_idx])
+            merged_text = merge_multi_strings(list(text_parts), overlap_ratio, last_overlap)
+            merged_conf = sum(confidences) / len(confidences)  # average confidence
+            remapped.append((merged_text, merged_conf))
+    return remapped

{onnxtr-0.6.3 → onnxtr-0.7.1}/onnxtr/models/recognition/predictor/base.py

@@ -36,7 +36,7 @@ class RecognitionPredictor(NestedObject):
         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.overlap_ratio = 0.5  # Ratio of overlap between neighboring crops
         self.target_ar = 6  # Target aspect ratio
 
     def __call__(
@@ -57,7 +57,7 @@ class RecognitionPredictor(NestedObject):
                 crops,  # type: ignore[arg-type]
                 self.critical_ar,
                 self.target_ar,
-                self.dil_factor,
+                self.overlap_ratio,
                 True,
             )
             if remapped:
@@ -74,6 +74,6 @@ class RecognitionPredictor(NestedObject):
 
         # Remap crops
         if self.split_wide_crops and remapped:
-            out = remap_preds(out, crop_map, self.dil_factor)
+            out = remap_preds(out, crop_map, self.overlap_ratio)
 
         return out

onnxtr-0.7.1/onnxtr/models/recognition/utils.py

@@ -0,0 +1,93 @@
+# Copyright (C) 2021-2025, 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 rapidfuzz.distance import Hamming
+
+__all__ = ["merge_strings", "merge_multi_strings"]
+
+
+def merge_strings(a: str, b: str, overlap_ratio: 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.
+        overlap_ratio: estimated ratio of overlapping characters.
+
+    Returns:
+        A merged character sequence.
+
+    Example::
+        >>> from doctr.models.recognition.utils import merge_strings
+        >>> merge_strings('abcd', 'cdefgh', 0.5)
+        'abcdefgh'
+        >>> merge_strings('abcdi', 'cdefgh', 0.5)
+        'abcdefgh'
+    """
+    seq_len = min(len(a), len(b))
+    if seq_len <= 1:  # One sequence is empty or will be after cropping in next step, return both to keep data
+        return a + b
+
+    a_crop, b_crop = a[:-1], b[1:]  # Remove last letter of "a" and first of "b", because they might be cut off
+    max_overlap = min(len(a_crop), len(b_crop))
+
+    # Compute Hamming distances for all possible overlaps
+    scores = [Hamming.distance(a_crop[-i:], b_crop[:i], processor=None) for i in range(1, max_overlap + 1)]
+
+    # Find zero-score matches
+    zero_matches = [i for i, score in enumerate(scores) if score == 0]
+
+    expected_overlap = round(len(b) * overlap_ratio) - 3  # adjust for cropping and index
+
+    # Case 1: One perfect match - exactly one zero score - just merge there
+    if len(zero_matches) == 1:
+        i = zero_matches[0]
+        return a_crop + b_crop[i + 1 :]
+
+    # Case 2: Multiple perfect matches - likely due to repeated characters.
+    # Use the estimated overlap length to choose the match closest to the expected alignment.
+    elif len(zero_matches) > 1:
+        best_i = min(zero_matches, key=lambda x: abs(x - expected_overlap))
+        return a_crop + b_crop[best_i + 1 :]
+
+    # Case 3: Absence of zero scores indicates that the same character in the image was recognized differently OR that
+    # the overlap was too small and we just need to merge the crops fully
+    if expected_overlap < -1:
+        return a + b
+    elif expected_overlap < 0:
+        return a_crop + b_crop
+
+    # Find best overlap by minimizing Hamming distance + distance from expected overlap size
+    combined_scores = [score + abs(i - expected_overlap) for i, score in enumerate(scores)]
+    best_i = combined_scores.index(min(combined_scores))
+    return a_crop + b_crop[best_i + 1 :]
+
+
+def merge_multi_strings(seq_list: list[str], overlap_ratio: float, last_overlap_ratio: float) -> str:
+    """
+    Merges consecutive string sequences with overlapping characters.
+
+    Args:
+        seq_list: list of sequences to merge. Sequences need to be ordered from left to right.
+        overlap_ratio: Estimated ratio of overlapping letters between neighboring strings.
+        last_overlap_ratio: Estimated ratio of overlapping letters for the last element in seq_list.
+
+    Returns:
+        A merged character sequence
+
+    Example::
+        >>> from doctr.models.recognition.utils import merge_multi_strings
+        >>> merge_multi_strings(['abc', 'bcdef', 'difghi', 'aijkl'], 0.5, 0.1)
+        'abcdefghijkl'
+    """
+    if not seq_list:
+        return ""
+    result = seq_list[0]
+    for i in range(1, len(seq_list)):
+        text_b = seq_list[i]
+        ratio = last_overlap_ratio if i == len(seq_list) - 1 else overlap_ratio
+        result = merge_strings(result, text_b, ratio)
+    return result

{onnxtr-0.6.3 → onnxtr-0.7.1}/onnxtr/models/recognition/zoo.py

@@ -22,6 +22,7 @@ ARCHS: list[str] = [
     "vitstr_small",
     "vitstr_base",
     "parseq",
+    "viptr_tiny",
 ]
 
 
@@ -35,7 +36,15 @@ def _predictor(
         _model = recognition.__dict__[arch](load_in_8_bit=load_in_8_bit, engine_cfg=engine_cfg)
     else:
         if not isinstance(
-            arch, (recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq)
+            arch,
+            (
+                recognition.CRNN,
+                recognition.SAR,
+                recognition.MASTER,
+                recognition.ViTSTR,
+                recognition.PARSeq,
+                recognition.VIPTR,
+            ),
         ):
             raise ValueError(f"unknown architecture: {type(arch)}")
         _model = arch