onnxtr 0.7.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

onnxtr/models/_utils.py

@@ -62,7 +62,7 @@ def estimate_orientation(
         thresh = img.astype(np.uint8)
 
     page_orientation, orientation_confidence = general_page_orientation or (None, 0.0)
-    if page_orientation and orientation_confidence >= min_confidence:
+    if page_orientation is not None and orientation_confidence >= min_confidence:
         # We rotate the image to the general orientation which improves the detection
         # No expand needed bitmap is already padded
         thresh = rotate_image(thresh, -page_orientation)
@@ -99,7 +99,7 @@ def estimate_orientation(
         estimated_angle = -round(median) if abs(median) != 0 else 0
 
     # combine with the general orientation and the estimated angle
-    if page_orientation and orientation_confidence >= min_confidence:
+    if page_orientation is not None and orientation_confidence >= min_confidence:
         # special case where the estimated angle is mostly wrong:
         # case 1: - and + swapped
         # case 2: estimated angle is completely wrong

onnxtr/models/detection/core.py

@@ -53,7 +53,7 @@ class DetectionPostProcessor(NestedObject):
 
         else:
             mask: np.ndarray = np.zeros((h, w), np.int32)
-            cv2.fillPoly(mask, [points.astype(np.int32)], 1.0)  # type: ignore[call-overload]
+            cv2.fillPoly(mask, [points.astype(np.int32)], 1.0)
             product = pred * mask
             return np.sum(product) / np.count_nonzero(product)
 

onnxtr/models/engine.py

@@ -5,13 +5,15 @@
 
 import logging
 import os
-from typing import Any
+from collections.abc import Callable
+from typing import Any, TypeAlias
 
 import numpy as np
 from onnxruntime import (
     ExecutionMode,
     GraphOptimizationLevel,
     InferenceSession,
+    RunOptions,
     SessionOptions,
     get_available_providers,
     get_device,
@@ -23,7 +25,9 @@ set_default_logger_severity(int(os.getenv("ORT_LOG_SEVERITY_LEVEL", 4)))
 from onnxtr.utils.data import download_from_url
 from onnxtr.utils.geometry import shape_translate
 
-__all__ = ["EngineConfig"]
+__all__ = ["EngineConfig", "RunOptionsProvider"]
+
+RunOptionsProvider: TypeAlias = Callable[[RunOptions], RunOptions]
 
 
 class EngineConfig:
@@ -38,9 +42,11 @@ class EngineConfig:
         self,
         providers: list[tuple[str, dict[str, Any]]] | list[str] | None = None,
         session_options: SessionOptions | None = None,
+        run_options_provider: RunOptionsProvider | None = None,
     ):
         self._providers = providers or self._init_providers()
         self._session_options = session_options or self._init_sess_opts()
+        self.run_options_provider = run_options_provider
 
     def _init_providers(self) -> list[tuple[str, dict[str, Any]]]:
         providers: Any = [("CPUExecutionProvider", {"arena_extend_strategy": "kSameAsRequested"})]
@@ -79,7 +85,7 @@ class EngineConfig:
         return self._session_options
 
     def __repr__(self) -> str:
-        return f"EngineConfig(providers={self.providers}"
+        return f"EngineConfig(providers={self.providers})"
 
 
 class Engine:
@@ -100,6 +106,7 @@ class Engine:
         self.model_path = archive_path
         self.session_options = engine_cfg.session_options
         self.providers = engine_cfg.providers
+        self.run_options_provider = engine_cfg.run_options_provider
         self.runtime = InferenceSession(archive_path, providers=self.providers, sess_options=self.session_options)
         self.runtime_inputs = self.runtime.get_inputs()[0]
         self.tf_exported = int(self.runtime_inputs.shape[-1]) == 3
@@ -109,6 +116,9 @@ class Engine:
         self.output_name = [output.name for output in self.runtime.get_outputs()]
 
     def run(self, inputs: np.ndarray) -> np.ndarray:
+        run_options = RunOptions()
+        if self.run_options_provider is not None:
+            run_options = self.run_options_provider(run_options)
         if self.tf_exported:
             inputs = shape_translate(inputs, format="BHWC")  # sanity check
         else:
@@ -117,8 +127,12 @@ class Engine:
             inputs = np.broadcast_to(inputs, (self.fixed_batch_size, *inputs.shape))
             # combine the results
             logits = np.concatenate(
-                [self.runtime.run(self.output_name, {self.runtime_inputs.name: batch})[0] for batch in inputs], axis=0
+                [
+                    self.runtime.run(self.output_name, {self.runtime_inputs.name: batch}, run_options=run_options)[0]
+                    for batch in inputs
+                ],
+                axis=0,
             )
         else:
-            logits = self.runtime.run(self.output_name, {self.runtime_inputs.name: inputs})[0]
+            logits = self.runtime.run(self.output_name, {self.runtime_inputs.name: inputs}, run_options=run_options)[0]
         return shape_translate(logits, format="BHWC")

onnxtr/models/recognition/models/crnn.py

@@ -22,9 +22,9 @@ default_cfgs: dict[str, dict[str, Any]] = {
         "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",
-        "url_8_bit": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.1.2/crnn_vgg16_bn_static_8_bit-bce050c7.onnx",
+        "vocab": VOCABS["french"],
+        "url": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.7.1/crnn_vgg16_bn-743599aa.onnx",
+        "url_8_bit": "https://github.com/felixdittrich92/OnnxTR/releases/download/v0.7.1/crnn_vgg16_bn_static_8_bit-df1b594d.onnx",
     },
     "crnn_mobilenet_v3_small": {
         "mean": (0.694, 0.695, 0.693),

onnxtr/utils/vocabs.py

@@ -264,8 +264,6 @@ VOCABS["estonian"] = VOCABS["english"] + "šžõäöüŠŽÕÄÖÜ"
 VOCABS["esperanto"] = re.sub(r"[QqWwXxYy]", "", VOCABS["english"]) + "ĉĝĥĵŝŭĈĜĤĴŜŬ" + "₷"
 
 VOCABS["french"] = VOCABS["english"] + "àâéèêëîïôùûüçÀÂÉÈÊËÎÏÔÙÛÜÇ"
-# NOTE: legacy french is outdated, but kept for compatibility
-VOCABS["legacy_french"] = VOCABS["latin"] + "°" + "àâéèêëîïôùûçÀÂÉÈËÎÏÔÙÛÇ" + _BASE_VOCABS["currency"]
 
 VOCABS["finnish"] = VOCABS["english"] + "äöÄÖ"
 

onnxtr/version.py

@@ -1 +1 @@
-__version__ = 'v0.7.1'
+__version__ = 'v0.8.0'

{onnxtr-0.7.1.dist-info → onnxtr-0.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxtr
-Version: 0.7.1
+Version: 0.8.0
 Summary: Onnx Text Recognition (OnnxTR): docTR Onnx-Wrapper for high-performance OCR on documents.
 Author-email: Felix Dittrich <felixdittrich92@gmail.com>
 Maintainer: Felix Dittrich
@@ -264,6 +264,7 @@ Requires-Dist: pytest>=5.3.2; extra == "testing"
 Requires-Dist: coverage[toml]>=4.5.4; extra == "testing"
 Requires-Dist: requests>=2.20.0; extra == "testing"
 Requires-Dist: pytest-memray>=1.7.0; extra == "testing"
+Requires-Dist: psutil>=7.0.0; extra == "testing"
 Provides-Extra: quality
 Requires-Dist: ruff>=0.1.5; extra == "quality"
 Requires-Dist: mypy>=0.812; extra == "quality"
@@ -278,6 +279,7 @@ Requires-Dist: pytest>=5.3.2; extra == "dev"
 Requires-Dist: coverage[toml]>=4.5.4; extra == "dev"
 Requires-Dist: requests>=2.20.0; extra == "dev"
 Requires-Dist: pytest-memray>=1.7.0; extra == "dev"
+Requires-Dist: psutil>=7.0.0; extra == "dev"
 Requires-Dist: ruff>=0.1.5; extra == "dev"
 Requires-Dist: mypy>=0.812; extra == "dev"
 Requires-Dist: pre-commit>=2.17.0; extra == "dev"
@@ -292,8 +294,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.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/)
+[![Socket Badge](https://socket.dev/api/badge/pypi/package/onnxtr/0.8.0?artifact_id=tar-gz)](https://socket.dev/pypi/package/onnxtr/overview/0.8.0/tar-gz)
+[![Pypi](https://img.shields.io/badge/pypi-v0.8.0-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)
@@ -480,6 +482,32 @@ predictor = ocr_predictor(
 )
 ```
 
+You can also dynamically configure whether the memory arena should shrink:
+
+```python
+from random import random
+from onnxruntime import RunOptions, SessionOptions
+
+from onnxtr.models import ocr_predictor, EngineConfig
+
+def arena_shrinkage_handler(run_options: RunOptions) -> RunOptions:
+  """
+  Shrink the memory arena on 10% of inference runs.
+  """
+  if random() < 0.1:
+    run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu:0")
+  return run_options
+
+engine_config = EngineConfig(run_options_provider=arena_shrinkage_handler)
+engine_config.session_options.enable_mem_pattern = False
+
+predictor = ocr_predictor(
+    det_engine_cfg=engine_config,
+    reco_engine_cfg=engine_config,
+    clf_engine_cfg=engine_config
+)
+```
+
 </details>
 
 ## Loading custom exported models

{onnxtr-0.7.1.dist-info → onnxtr-0.8.0.dist-info}/RECORD

@@ -1,7 +1,7 @@
 onnxtr/__init__.py,sha256=h7Wc2tuHLsaoCk5xNpEFEK-g11A6SJA7nAasA76TQ_Y,100
 onnxtr/file_utils.py,sha256=rlflwU-KTHQ6wD1GU2EqPJs4y7gTyZSUnhtx9w2lIOk,1067
 onnxtr/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-onnxtr/version.py,sha256=XMmaM5zJmM1OR4hLdJ8j4e4PQGtWFhE8q9jYdfTRB4c,23
+onnxtr/version.py,sha256=XbvKmQHvfaLsOHsQx8r-8Tg2eqiGheb-FcQ7cNGg5ro,23
 onnxtr/contrib/__init__.py,sha256=0Jo052I1DIdbNVYQBg213UJBNH3UKZIOcn2HVusQlj4,39
 onnxtr/contrib/artefacts.py,sha256=4uJyXrl2YEjeWaGC8pAVC5227eVzb34vTSbAe0fzGbw,5323
 onnxtr/contrib/base.py,sha256=SPPWd8SVwZr4kGPI-yM8UdQguEZ0awRBUZGjGKZEC-U,3135
@@ -12,9 +12,9 @@ onnxtr/io/image.py,sha256=ehOcXzH1eFjwLO3pwiV3bUwGcn4SdW-FO9DVbUjWw2M,1700
 onnxtr/io/pdf.py,sha256=cO_e9MWh2uCu20a9eLqp3y3GS9idpROkOuNkCU5yE8U,1327
 onnxtr/io/reader.py,sha256=JBDoGXaQsiV3eZxDu4Sd8aeULIEmLV5oNGqGXLAmtBk,2755
 onnxtr/models/__init__.py,sha256=QTfZlqUyv1d7NUCbGIUFM1DLOOXe-cqHZ7uaKkGdXvk,157
-onnxtr/models/_utils.py,sha256=m_6Qnx9QNIxmf_rrYmjd0htFmGa6ySZLoQ9kPjtMlpA,6424
+onnxtr/models/_utils.py,sha256=IqXlismua45W3Uo451oVNrnMnUCrxKoRqKJ5TKyypYE,6448
 onnxtr/models/builder.py,sha256=aaMHLZl8fayeWsXHsvSy5Tq2yQWTO0gpj4Ob4kNZ0xk,14076
-onnxtr/models/engine.py,sha256=mtU4kKoWcgkcLELIJKTxYD2GMFAVZ1b_i7tSG2VeYQ8,5113
+onnxtr/models/engine.py,sha256=FafLMq_qtNqyQSPO-pvYiaX9BM5YINENhQdp02EWpu0,5732
 onnxtr/models/zoo.py,sha256=44H7NfE4DWH7KgmF8vod9fMlPnxw41dwjpwjuIJkLHg,5303
 onnxtr/models/classification/__init__.py,sha256=h1bZs55iLJBMATtzS4ntTKwfD6OGXBiiqGv_hEnOFnE,41
 onnxtr/models/classification/zoo.py,sha256=4hD51Mdf5gkSoQ4g9zTBvurbBdrsaFFEl5WFpptBtNk,4271
@@ -23,7 +23,7 @@ onnxtr/models/classification/models/mobilenet.py,sha256=pRI97DvBpBE-h2VymUG6sw_J
 onnxtr/models/classification/predictor/__init__.py,sha256=ERmmOxz_9mUkIuccNbzUa5Y6gVLLVDdyc4cCxbCCUbY,20
 onnxtr/models/classification/predictor/base.py,sha256=hpLYmhn8PCbx46zYBM8rZiy2vabpN3YkviwmmbVaqRo,2313
 onnxtr/models/detection/__init__.py,sha256=h1bZs55iLJBMATtzS4ntTKwfD6OGXBiiqGv_hEnOFnE,41
-onnxtr/models/detection/core.py,sha256=8avDAHBh-S3jjQ-4WTcbXQvUgwrvJtFPti2ghUtE9uA,3497
+onnxtr/models/detection/core.py,sha256=0oTo-NxtsvVwXqCpNHv98uiXoOcUk8WsIRi5X65bnc8,3466
 onnxtr/models/detection/zoo.py,sha256=RbzG9nR0Oa9bV0vwEBer2GDJMg1XvthpRLBxD2MsUvQ,3385
 onnxtr/models/detection/_utils/__init__.py,sha256=oPkIYbySSbLsOk02wVPNO9bUuywC47YjaenfyTwfOsw,20
 onnxtr/models/detection/_utils/base.py,sha256=MUn3cHUSq-dIKAA021ZUrTm0SLpAC9qsiw1BAWvcppc,2300
@@ -47,7 +47,7 @@ onnxtr/models/recognition/core.py,sha256=V6DMa3ABT0amr9vPSQSyo0dZqK-Y8qmA__XWhOx
 onnxtr/models/recognition/utils.py,sha256=jsKZl7XbGKCgEhdH4xYDdf9xrrDR0vS3RqbwGIg17z0,3756
 onnxtr/models/recognition/zoo.py,sha256=nNDf60Et6wE60cJhcqwxNbCPQfcv2pgQIK6KeGVH0WE,3018
 onnxtr/models/recognition/models/__init__.py,sha256=hbjiDSCgh_EUEM7q9aSbcK7ffszMvsDWW89-MVN21Vk,126
-onnxtr/models/recognition/models/crnn.py,sha256=9NwYTsrAalNHuNkxURe574v4tvKY-BsJOb25X4XDIKk,8786
+onnxtr/models/recognition/models/crnn.py,sha256=fZDrHp6ojs_-r1VZm9J8jMiRQ4lA760pBZXTzzi0Jrk,8779
 onnxtr/models/recognition/models/master.py,sha256=H_KcoJCK8nElyID5FKJdld5xapu5ZOWVcH5_H68_N1I,4669
 onnxtr/models/recognition/models/parseq.py,sha256=qe7lavwFT37SPLSMWoPaoCYo7yrTYGcCYqQraTwIuQY,4512
 onnxtr/models/recognition/models/sar.py,sha256=5v0VOPACvB_-I_kDh57pA00GmunIh6frD_4I6ojSjiQ,4523
@@ -67,10 +67,10 @@ onnxtr/utils/multithreading.py,sha256=P8LXpaZSKRYsYZw6SRlWWXiQMDswk8zEgDSbCYVHzS
 onnxtr/utils/reconstitution.py,sha256=2SUkaReD5eJ5NFs7V7Z0m4N4lUr8QB7oH1mSHTC3vLA,6119
 onnxtr/utils/repr.py,sha256=lAcKRShXLpJ2DvsS2qf3vYq0lJrMJa30kLYXmfN3rV0,2105
 onnxtr/utils/visualization.py,sha256=MImlxVKOLXPriJYSrBfjZfJw0GEZ1gT8eni3KXZHrXk,9831
-onnxtr/utils/vocabs.py,sha256=NI96-ZvjpfyKfdbyaBo7bcUCmJdamamsQbE0WBdUFyk,99238
-onnxtr-0.7.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-onnxtr-0.7.1.dist-info/METADATA,sha256=MV34cdklkhwt6FsyYiTONjFHhVwm7Xe9jNJt5f7dgKo,35072
-onnxtr-0.7.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnxtr-0.7.1.dist-info/top_level.txt,sha256=r_MSUTpspp4pWEEWvly-s7ZkfCg1KwrK6-kBlXkWKU8,7
-onnxtr-0.7.1.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-onnxtr-0.7.1.dist-info/RECORD,,
+onnxtr/utils/vocabs.py,sha256=i7LIv6tUiTXMcqHZfDgU43knJ-9zBptF6loXdft4KY4,99049
+onnxtr-0.8.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+onnxtr-0.8.0.dist-info/METADATA,sha256=JLbYn1isqfSmHlTpKVpTxFmx0HW1YQpxEeVPv7ID29s,35915
+onnxtr-0.8.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnxtr-0.8.0.dist-info/top_level.txt,sha256=r_MSUTpspp4pWEEWvly-s7ZkfCg1KwrK6-kBlXkWKU8,7
+onnxtr-0.8.0.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+onnxtr-0.8.0.dist-info/RECORD,,