onnxtr 0.3.1__tar.gz → 0.4.0__tar.gz

{onnxtr-0.3.1 → onnxtr-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxtr
-Version: 0.3.1
+Version: 0.4.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
@@ -233,6 +233,7 @@ Requires-Dist: pyclipper<2.0.0,>=1.2.0
 Requires-Dist: shapely<3.0.0,>=1.6.0
 Requires-Dist: rapidfuzz<4.0.0,>=3.0.0
 Requires-Dist: langdetect<2.0.0,>=1.0.9
+Requires-Dist: huggingface-hub<1.0.0,>=0.23.0
 Requires-Dist: Pillow>=9.2.0
 Requires-Dist: defusedxml>=0.7.0
 Requires-Dist: anyascii>=0.3.2
@@ -275,7 +276,7 @@ Requires-Dist: pre-commit>=2.17.0; extra == "dev"
 [![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)
-[![Pypi](https://img.shields.io/badge/pypi-v0.3.1-blue.svg)](https://pypi.org/project/OnnxTR/)
+[![Pypi](https://img.shields.io/badge/pypi-v0.3.2-blue.svg)](https://pypi.org/project/OnnxTR/)
 
 > :warning: Please note that this is a wrapper around the [doctr](https://github.com/mindee/doctr) library to provide a Onnx pipeline for docTR. For feature requests, which are not directly related to the Onnx pipeline, please refer to the base project.
 
@@ -449,6 +450,69 @@ det_model = linknet_resnet18("path_to_custom_model.onnx")
 model = ocr_predictor(det_arch=det_model, reco_arch=reco_model)
 ```
 
+## Loading models from HuggingFace Hub
+
+You can also load models from the HuggingFace Hub:
+
+```python
+from onnxtr.io import DocumentFile
+from onnxtr.models import ocr_predictor, from_hub
+
+img = DocumentFile.from_images(['<image_path>'])
+# Load your model from the hub
+model = from_hub('onnxtr/my-model')
+
+# Pass it to the predictor
+# If your model is a recognition model:
+predictor = ocr_predictor(
+    det_arch='db_mobilenet_v3_large',
+    reco_arch=model
+)
+
+# If your model is a detection model:
+predictor = ocr_predictor(
+    det_arch=model,
+    reco_arch='crnn_mobilenet_v3_small'
+)
+
+# Get your predictions
+res = predictor(img)
+```
+
+HF Hub search: [here](https://huggingface.co/models?search=onnxtr).
+
+Collection: [here](https://huggingface.co/collections/Felix92/onnxtr-66bf213a9f88f7346c90e842)
+
+Or push your own models to the hub:
+
+```python
+from onnxtr.models import parseq, push_to_hf_hub, login_to_hub
+from onnxtr.utils.vocabs import VOCABS
+
+# Login to the hub
+login_to_hub()
+
+# Recogniton model
+model = parseq("~/onnxtr-parseq-multilingual-v1.onnx", vocab=VOCABS["multilingual"])
+push_to_hf_hub(
+    model,
+    model_name="onnxtr-parseq-multilingual-v1",
+    task="recognition",  # The task for which the model is intended [detection, recognition, classification]
+    arch="parseq",  # The name of the model architecture
+    override=False  # Set to `True` if you want to override an existing model / repository
+)
+
+# Detection model
+model = linknet_resnet18("~/onnxtr-linknet-resnet18.onnx")
+push_to_hf_hub(
+    model,
+    model_name="onnxtr-linknet-resnet18",
+    task="detection",
+    arch="linknet_resnet18",
+    override=True
+)
+```
+
 ## Models architectures
 
 Credits where it's due: this repository provides ONNX models for the following architectures, converted from the docTR models:

{onnxtr-0.3.1 → onnxtr-0.4.0}/README.md

@@ -7,7 +7,7 @@
 [![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)
-[![Pypi](https://img.shields.io/badge/pypi-v0.3.1-blue.svg)](https://pypi.org/project/OnnxTR/)
+[![Pypi](https://img.shields.io/badge/pypi-v0.3.2-blue.svg)](https://pypi.org/project/OnnxTR/)
 
 > :warning: Please note that this is a wrapper around the [doctr](https://github.com/mindee/doctr) library to provide a Onnx pipeline for docTR. For feature requests, which are not directly related to the Onnx pipeline, please refer to the base project.
 
@@ -181,6 +181,69 @@ det_model = linknet_resnet18("path_to_custom_model.onnx")
 model = ocr_predictor(det_arch=det_model, reco_arch=reco_model)
 ```
 
+## Loading models from HuggingFace Hub
+
+You can also load models from the HuggingFace Hub:
+
+```python
+from onnxtr.io import DocumentFile
+from onnxtr.models import ocr_predictor, from_hub
+
+img = DocumentFile.from_images(['<image_path>'])
+# Load your model from the hub
+model = from_hub('onnxtr/my-model')
+
+# Pass it to the predictor
+# If your model is a recognition model:
+predictor = ocr_predictor(
+    det_arch='db_mobilenet_v3_large',
+    reco_arch=model
+)
+
+# If your model is a detection model:
+predictor = ocr_predictor(
+    det_arch=model,
+    reco_arch='crnn_mobilenet_v3_small'
+)
+
+# Get your predictions
+res = predictor(img)
+```
+
+HF Hub search: [here](https://huggingface.co/models?search=onnxtr).
+
+Collection: [here](https://huggingface.co/collections/Felix92/onnxtr-66bf213a9f88f7346c90e842)
+
+Or push your own models to the hub:
+
+```python
+from onnxtr.models import parseq, push_to_hf_hub, login_to_hub
+from onnxtr.utils.vocabs import VOCABS
+
+# Login to the hub
+login_to_hub()
+
+# Recogniton model
+model = parseq("~/onnxtr-parseq-multilingual-v1.onnx", vocab=VOCABS["multilingual"])
+push_to_hf_hub(
+    model,
+    model_name="onnxtr-parseq-multilingual-v1",
+    task="recognition",  # The task for which the model is intended [detection, recognition, classification]
+    arch="parseq",  # The name of the model architecture
+    override=False  # Set to `True` if you want to override an existing model / repository
+)
+
+# Detection model
+model = linknet_resnet18("~/onnxtr-linknet-resnet18.onnx")
+push_to_hf_hub(
+    model,
+    model_name="onnxtr-linknet-resnet18",
+    task="detection",
+    arch="linknet_resnet18",
+    override=True
+)
+```
+
 ## Models architectures
 
 Credits where it's due: this repository provides ONNX models for the following architectures, converted from the docTR models:

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/__init__.py

@@ -3,3 +3,4 @@ from .classification import *
 from .detection import *
 from .recognition import *
 from .zoo import *
+from .factory import *

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/classification/models/mobilenet.py

@@ -56,6 +56,7 @@ class MobileNetV3(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.cfg = cfg
 
     def __call__(

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/detection/models/differentiable_binarization.py

@@ -64,8 +64,10 @@ class DBNet(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.cfg = cfg
         self.assume_straight_pages = assume_straight_pages
+
         self.postprocessor = GeneralDetectionPostProcessor(
             assume_straight_pages=self.assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh
         )

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/detection/models/fast.py

@@ -62,6 +62,7 @@ class FAST(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.cfg = cfg
         self.assume_straight_pages = assume_straight_pages
 

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/detection/models/linknet.py

@@ -64,6 +64,7 @@ class LinkNet(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.cfg = cfg
         self.assume_straight_pages = assume_straight_pages
 

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/engine.py

@@ -90,6 +90,8 @@ class Engine:
     def __init__(self, url: str, engine_cfg: Optional[EngineConfig] = None, **kwargs: Any) -> None:
         engine_cfg = engine_cfg if isinstance(engine_cfg, EngineConfig) else EngineConfig()
         archive_path = download_from_url(url, cache_subdir="models", **kwargs) if "http" in url else url
+        # Store model path for each model
+        self.model_path = archive_path
         self.session_options = engine_cfg.session_options
         self.providers = engine_cfg.providers
         self.runtime = InferenceSession(archive_path, providers=self.providers, sess_options=self.session_options)

onnxtr-0.4.0/onnxtr/models/factory/__init__.py

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

onnxtr-0.4.0/onnxtr/models/factory/hub.py

@@ -0,0 +1,224 @@
+# 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.
+
+# Inspired by: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/hub.py
+
+import json
+import logging
+import os
+import shutil
+import subprocess
+import textwrap
+from pathlib import Path
+from typing import Any, Optional
+
+from huggingface_hub import (
+    HfApi,
+    Repository,
+    get_token,
+    get_token_permission,
+    hf_hub_download,
+    login,
+)
+
+from onnxtr import models
+from onnxtr.models.engine import EngineConfig
+
+__all__ = ["login_to_hub", "push_to_hf_hub", "from_hub", "_save_model_and_config_for_hf_hub"]
+
+
+AVAILABLE_ARCHS = {
+    "classification": models.classification.zoo.ORIENTATION_ARCHS,
+    "detection": models.detection.zoo.ARCHS,
+    "recognition": models.recognition.zoo.ARCHS,
+}
+
+
+def login_to_hub() -> None:  # pragma: no cover
+    """Login to huggingface hub"""
+    access_token = get_token()
+    if access_token is not None and get_token_permission(access_token):
+        logging.info("Huggingface Hub token found and valid")
+        login(token=access_token, write_permission=True)
+    else:
+        login()
+    # check if git lfs is installed
+    try:
+        subprocess.call(["git", "lfs", "version"])
+    except FileNotFoundError:
+        raise OSError(
+            "Looks like you do not have git-lfs installed, please install. \
+                      You can install from https://git-lfs.github.com/. \
+                      Then run `git lfs install` (you only have to do this once)."
+        )
+
+
+def _save_model_and_config_for_hf_hub(model: Any, save_dir: str, arch: str, task: str) -> None:
+    """Save model and config to disk for pushing to huggingface hub
+
+    Args:
+    ----
+        model: Onnx model to be saved
+        save_dir: directory to save model and config
+        arch: architecture name
+        task: task name
+    """
+    save_directory = Path(save_dir)
+    shutil.copy2(model.model_path, save_directory / "model.onnx")
+
+    config_path = save_directory / "config.json"
+
+    # add model configuration
+    model_config = model.cfg
+    model_config["arch"] = arch
+    model_config["task"] = task
+
+    with config_path.open("w") as f:
+        json.dump(model_config, f, indent=2, ensure_ascii=False)
+
+
+def push_to_hf_hub(
+    model: Any, model_name: str, task: str, override: bool = False, **kwargs
+) -> None:  # pragma: no cover
+    """Save model and its configuration on HF hub
+
+    >>> from onnxtr.models import login_to_hub, push_to_hf_hub
+    >>> from onnxtr.models.recognition import crnn_mobilenet_v3_small
+    >>> login_to_hub()
+    >>> model = crnn_mobilenet_v3_small()
+    >>> push_to_hf_hub(model, 'my-model', 'recognition', arch='crnn_mobilenet_v3_small')
+
+    Args:
+    ----
+        model: Onnx model to be saved
+        model_name: name of the model which is also the repository name
+        task: task name
+        override: whether to override the existing model / repo on HF hub
+        **kwargs: keyword arguments for push_to_hf_hub
+    """
+    run_config = kwargs.get("run_config", None)
+    arch = kwargs.get("arch", None)
+
+    if run_config is None and arch is None:
+        raise ValueError("run_config or arch must be specified")
+    if task not in ["classification", "detection", "recognition"]:
+        raise ValueError("task must be one of classification, detection, recognition")
+
+    # default readme
+    readme = textwrap.dedent(
+        f"""
+    ---
+    language:
+    - en
+    - fr
+    license: apache-2.0
+    ---
+
+    <p align="center">
+    <img src="https://github.com/felixdittrich92/OnnxTR/raw/main/docs/images/logo.jpg" width="40%">
+    </p>
+
+    **Optical Character Recognition made seamless & accessible to anyone, powered by Onnxruntime**
+
+    ## Task: {task}
+
+    https://github.com/felixdittrich92/OnnxTR
+
+    ### Example usage:
+
+    ```python
+    >>> from onnxtr.io import DocumentFile
+    >>> from onnxtr.models import ocr_predictor, from_hub
+
+    >>> img = DocumentFile.from_images(['<image_path>'])
+    >>> # Load your model from the hub
+    >>> model = from_hub('onnxtr/my-model')
+
+    >>> # Pass it to the predictor
+    >>> # If your model is a recognition model:
+    >>> predictor = ocr_predictor(det_arch='db_mobilenet_v3_large',
+    >>>                           reco_arch=model)
+
+    >>> # If your model is a detection model:
+    >>> predictor = ocr_predictor(det_arch=model,
+    >>>                           reco_arch='crnn_mobilenet_v3_small')
+
+    >>> # Get your predictions
+    >>> res = predictor(img)
+    ```
+    """
+    )
+
+    # add run configuration to readme if available
+    if run_config is not None:
+        arch = run_config.arch
+        readme += textwrap.dedent(
+            f"""### Run Configuration
+                                  \n{json.dumps(vars(run_config), indent=2, ensure_ascii=False)}"""
+        )
+
+    if arch not in AVAILABLE_ARCHS[task]:
+        raise ValueError(
+            f"Architecture: {arch} for task: {task} not found.\
+                         \nAvailable architectures: {AVAILABLE_ARCHS}"
+        )
+
+    commit_message = f"Add {model_name} model"
+
+    local_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub", model_name)
+    repo_url = HfApi().create_repo(model_name, token=get_token(), exist_ok=override)
+    repo = Repository(local_dir=local_cache_dir, clone_from=repo_url)
+
+    with repo.commit(commit_message):
+        _save_model_and_config_for_hf_hub(model, repo.local_dir, arch=arch, task=task)
+        readme_path = Path(repo.local_dir) / "README.md"
+        readme_path.write_text(readme)
+
+    repo.git_push()
+
+
+def from_hub(repo_id: str, engine_cfg: Optional[EngineConfig] = None, **kwargs: Any):
+    """Instantiate & load a pretrained model from HF hub.
+
+    >>> from onnxtr.models import from_hub
+    >>> model = from_hub("onnxtr/my-model")
+
+    Args:
+    ----
+        repo_id: HuggingFace model hub repo
+        engine_cfg: configuration for the inference engine (optional)
+        kwargs: kwargs of `hf_hub_download`
+
+    Returns:
+    -------
+        Model loaded with the checkpoint
+    """
+    # Get the config
+    with open(hf_hub_download(repo_id, filename="config.json", **kwargs), "rb") as f:
+        cfg = json.load(f)
+        model_path = hf_hub_download(repo_id, filename="model.onnx", **kwargs)
+
+    arch = cfg["arch"]
+    task = cfg["task"]
+    cfg.pop("arch")
+    cfg.pop("task")
+
+    if task == "classification":
+        model = models.classification.__dict__[arch](model_path, classes=cfg["classes"], engine_cfg=engine_cfg)
+    elif task == "detection":
+        model = models.detection.__dict__[arch](model_path, engine_cfg=engine_cfg)
+    elif task == "recognition":
+        model = models.recognition.__dict__[arch](
+            model_path, input_shape=cfg["input_shape"], vocab=cfg["vocab"], engine_cfg=engine_cfg
+        )
+
+    # convert all values which are lists to tuples
+    for key, value in cfg.items():
+        if isinstance(value, list):
+            cfg[key] = tuple(value)
+    # update model cfg
+    model.cfg = cfg
+
+    return model

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/recognition/models/crnn.py

@@ -129,8 +129,10 @@ class CRNN(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.vocab = vocab
         self.cfg = cfg
+
         self.postprocessor = CRNNPostProcessor(self.vocab)
 
     def __call__(

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/recognition/models/master.py

@@ -53,6 +53,7 @@ class MASTER(Engine):
 
         self.vocab = vocab
         self.cfg = cfg
+
         self.postprocessor = MASTERPostProcessor(vocab=self.vocab)
 
     def __call__(

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/recognition/models/parseq.py

@@ -49,8 +49,10 @@ class PARSeq(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.vocab = vocab
         self.cfg = cfg
+
         self.postprocessor = PARSeqPostProcessor(vocab=self.vocab)
 
     def __call__(

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/recognition/models/sar.py

@@ -49,8 +49,10 @@ class SAR(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.vocab = vocab
         self.cfg = cfg
+
         self.postprocessor = SARPostProcessor(self.vocab)
 
     def __call__(

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/recognition/models/vitstr.py

@@ -57,6 +57,7 @@ class ViTSTR(Engine):
         **kwargs: Any,
     ) -> None:
         super().__init__(url=model_path, engine_cfg=engine_cfg, **kwargs)
+
         self.vocab = vocab
         self.cfg = cfg
 

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/transforms/base.py

@@ -5,8 +5,8 @@
 
 from typing import Tuple, Union
 
-import cv2
 import numpy as np
+from PIL import Image, ImageOps
 
 __all__ = ["Resize", "Normalize"]
 
@@ -17,64 +17,51 @@ class Resize:
     def __init__(
         self,
         size: Union[int, Tuple[int, int]],
-        interpolation=cv2.INTER_LINEAR,
+        interpolation=Image.Resampling.BILINEAR,
         preserve_aspect_ratio: bool = False,
         symmetric_pad: bool = False,
     ) -> None:
-        super().__init__()
-        self.size = size
+        self.size = size if isinstance(size, tuple) else (size, size)
         self.interpolation = interpolation
         self.preserve_aspect_ratio = preserve_aspect_ratio
         self.symmetric_pad = symmetric_pad
         self.output_size = size if isinstance(size, tuple) else (size, size)
 
-        if not isinstance(self.size, (int, tuple, list)):
-            raise AssertionError("size should be either a tuple, a list or an int")
+        if not isinstance(self.size, (tuple, int)):
+            raise AssertionError("size should be either a tuple or an int")
 
-    def __call__(
-        self,
-        img: np.ndarray,
-    ) -> np.ndarray:
-        if img.ndim == 3:
-            h, w = img.shape[0:2]
-        else:
-            h, w = img.shape[1:3]
-        sh, sw = self.size if isinstance(self.size, tuple) else (self.size, self.size)
+    def __call__(self, img: np.ndarray) -> np.ndarray:
+        img = (img * 255).astype(np.uint8) if img.dtype != np.uint8 else img
+        h, w = img.shape[:2] if img.ndim == 3 else img.shape[1:3]
+        sh, sw = self.size
 
-        # Calculate aspect ratio of the image
-        aspect = w / h
+        if not self.preserve_aspect_ratio:
+            return np.array(Image.fromarray(img).resize((sw, sh), resample=self.interpolation))
 
-        # Compute scaling and padding sizes
-        if self.preserve_aspect_ratio:
-            if aspect > 1:  # Horizontal image
-                new_w = sw
-                new_h = int(sw / aspect)
-            elif aspect < 1:  # Vertical image
-                new_h = sh
-                new_w = int(sh * aspect)
-            else:  # Square image
-                new_h, new_w = sh, sw
-
-            img_resized = cv2.resize(img, (new_w, new_h), interpolation=self.interpolation)
-
-            # Calculate padding
-            pad_top = max((sh - new_h) // 2, 0)
-            pad_bottom = max(sh - new_h - pad_top, 0)
-            pad_left = max((sw - new_w) // 2, 0)
-            pad_right = max(sw - new_w - pad_left, 0)
-
-            # Pad the image
-            img_resized = cv2.copyMakeBorder(  # type: ignore[call-overload]
-                img_resized, pad_top, pad_bottom, pad_left, pad_right, borderType=cv2.BORDER_CONSTANT, value=0
-            )
-
-            # Ensure the image matches the target size by resizing it again if needed
-            img_resized = cv2.resize(img_resized, (sw, sh), interpolation=self.interpolation)
+        actual_ratio = h / w
+        target_ratio = sh / sw
+
+        if target_ratio == actual_ratio:
+            return np.array(Image.fromarray(img).resize((sw, sh), resample=self.interpolation))
+
+        if actual_ratio > target_ratio:
+            tmp_size = (int(sh / actual_ratio), sh)
         else:
-            # Resize the image without preserving aspect ratio
-            img_resized = cv2.resize(img, (sw, sh), interpolation=self.interpolation)
+            tmp_size = (sw, int(sw * actual_ratio))
+
+        img_resized = Image.fromarray(img).resize(tmp_size, resample=self.interpolation)
+        pad_left = pad_top = 0
+        pad_right = sw - img_resized.width
+        pad_bottom = sh - img_resized.height
+
+        if self.symmetric_pad:
+            pad_left = pad_right // 2
+            pad_right -= pad_left
+            pad_top = pad_bottom // 2
+            pad_bottom -= pad_top
 
-        return img_resized
+        img_resized = ImageOps.expand(img_resized, (pad_left, pad_top, pad_right, pad_bottom))
+        return np.array(img_resized)
 
     def __repr__(self) -> str:
         interpolate_str = self.interpolation

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/utils/fonts.py

@@ -5,14 +5,16 @@
 
 import logging
 import platform
-from typing import Optional
+from typing import Optional, Union
 
 from PIL import ImageFont
 
 __all__ = ["get_font"]
 
 
-def get_font(font_family: Optional[str] = None, font_size: int = 13) -> ImageFont.ImageFont:
+def get_font(
+    font_family: Optional[str] = None, font_size: int = 13
+) -> Union[ImageFont.FreeTypeFont, ImageFont.ImageFont]:
     """Resolves a compatible ImageFont for the system
 
     Args:
@@ -29,7 +31,7 @@ def get_font(font_family: Optional[str] = None, font_size: int = 13) -> ImageFon
         try:
             font = ImageFont.truetype("FreeMono.ttf" if platform.system() == "Linux" else "Arial.ttf", font_size)
         except OSError:  # pragma: no cover
-            font = ImageFont.load_default()
+            font = ImageFont.load_default()  # type: ignore[assignment]
             logging.warning(
                 "unable to load recommended font family. Loading default PIL font,"
                 "font size issues may be expected."

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/utils/vocabs.py

@@ -17,9 +17,14 @@ VOCABS: Dict[str, str] = {
     "ancient_greek": "αβγδεζηθικλμνξοπρστυφχψωΑΒΓΔΕΖΗΘΙΚΛΜΝΞΟΠΡΣΤΥΦΧΨΩ",
     "arabic_letters": "ءآأؤإئابةتثجحخدذرزسشصضطظعغـفقكلمنهوىي",
     "persian_letters": "پچڢڤگ",
-    "hindi_digits": "٠١٢٣٤٥٦٧٨٩",
+    "arabic_digits": "٠١٢٣٤٥٦٧٨٩",
     "arabic_diacritics": "ًٌٍَُِّْ",
     "arabic_punctuation": "؟؛«»—",
+    "hindi_letters": "अआइईउऊऋॠऌॡएऐओऔअंअःकखगघङचछजझञटठडढणतथदधनपफबभमयरलवशषसह",
+    "hindi_digits": "०१२३४५६७८९",
+    "hindi_punctuation": "।,?!:्ॐ॰॥॰",
+    "bangla_letters": "অআইঈউঊঋএঐওঔকখগঘঙচছজঝঞটঠডঢণতথদধনপফবভমযরলশষসহ়ঽািীুূৃেৈোৌ্ৎংঃঁ",
+    "bangla_digits": "০১২৩৪৫৬৭৮৯",
 }
 
 VOCABS["latin"] = VOCABS["digits"] + VOCABS["ascii_letters"] + VOCABS["punctuation"]
@@ -32,7 +37,7 @@ VOCABS["italian"] = VOCABS["english"] + "àèéìíîòóùúÀÈÉÌÍÎÒÓÙ
 VOCABS["german"] = VOCABS["english"] + "äöüßÄÖÜẞ"
 VOCABS["arabic"] = (
     VOCABS["digits"]
-    + VOCABS["hindi_digits"]
+    + VOCABS["arabic_digits"]
     + VOCABS["arabic_letters"]
     + VOCABS["persian_letters"]
     + VOCABS["arabic_diacritics"]
@@ -48,10 +53,12 @@ VOCABS["finnish"] = VOCABS["english"] + "äöÄÖ"
 VOCABS["swedish"] = VOCABS["english"] + "åäöÅÄÖ"
 VOCABS["vietnamese"] = (
     VOCABS["english"]
-    + "áàảạãăắằẳẵặâấầẩẫậéèẻẽẹêếềểễệóòỏõọôốồổộỗơớờởợỡúùủũụưứừửữựiíìỉĩịýỳỷỹỵ"
-    + "ÁÀẢẠÃĂẮẰẲẴẶÂẤẦẨẪẬÉÈẺẼẸÊẾỀỂỄỆÓÒỎÕỌÔỐỒỔỘỖƠỚỜỞỢỠÚÙỦŨỤƯỨỪỬỮỰIÍÌỈĨỊÝỲỶỸỴ"
+    + "áàảạãăắằẳẵặâấầẩẫậđéèẻẽẹêếềểễệóòỏõọôốồổộỗơớờởợỡúùủũụưứừửữựiíìỉĩịýỳỷỹỵ"
+    + "ÁÀẢẠÃĂẮẰẲẴẶÂẤẦẨẪẬĐÉÈẺẼẸÊẾỀỂỄỆÓÒỎÕỌÔỐỒỔỘỖƠỚỜỞỢỠÚÙỦŨỤƯỨỪỬỮỰIÍÌỈĨỊÝỲỶỸỴ"
 )
 VOCABS["hebrew"] = VOCABS["english"] + "אבגדהוזחטיכלמנסעפצקרשת" + "₪"
+VOCABS["hindi"] = VOCABS["hindi_letters"] + VOCABS["hindi_digits"] + VOCABS["hindi_punctuation"]
+VOCABS["bangla"] = VOCABS["bangla_letters"] + VOCABS["bangla_digits"]
 VOCABS["multilingual"] = "".join(
     dict.fromkeys(
         VOCABS["french"]

onnxtr-0.4.0/onnxtr/version.py

@@ -0,0 +1 @@
+__version__ = 'v0.4.0'

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxtr
-Version: 0.3.1
+Version: 0.4.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
@@ -233,6 +233,7 @@ Requires-Dist: pyclipper<2.0.0,>=1.2.0
 Requires-Dist: shapely<3.0.0,>=1.6.0
 Requires-Dist: rapidfuzz<4.0.0,>=3.0.0
 Requires-Dist: langdetect<2.0.0,>=1.0.9
+Requires-Dist: huggingface-hub<1.0.0,>=0.23.0
 Requires-Dist: Pillow>=9.2.0
 Requires-Dist: defusedxml>=0.7.0
 Requires-Dist: anyascii>=0.3.2
@@ -275,7 +276,7 @@ Requires-Dist: pre-commit>=2.17.0; extra == "dev"
 [![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)
-[![Pypi](https://img.shields.io/badge/pypi-v0.3.1-blue.svg)](https://pypi.org/project/OnnxTR/)
+[![Pypi](https://img.shields.io/badge/pypi-v0.3.2-blue.svg)](https://pypi.org/project/OnnxTR/)
 
 > :warning: Please note that this is a wrapper around the [doctr](https://github.com/mindee/doctr) library to provide a Onnx pipeline for docTR. For feature requests, which are not directly related to the Onnx pipeline, please refer to the base project.
 
@@ -449,6 +450,69 @@ det_model = linknet_resnet18("path_to_custom_model.onnx")
 model = ocr_predictor(det_arch=det_model, reco_arch=reco_model)
 ```
 
+## Loading models from HuggingFace Hub
+
+You can also load models from the HuggingFace Hub:
+
+```python
+from onnxtr.io import DocumentFile
+from onnxtr.models import ocr_predictor, from_hub
+
+img = DocumentFile.from_images(['<image_path>'])
+# Load your model from the hub
+model = from_hub('onnxtr/my-model')
+
+# Pass it to the predictor
+# If your model is a recognition model:
+predictor = ocr_predictor(
+    det_arch='db_mobilenet_v3_large',
+    reco_arch=model
+)
+
+# If your model is a detection model:
+predictor = ocr_predictor(
+    det_arch=model,
+    reco_arch='crnn_mobilenet_v3_small'
+)
+
+# Get your predictions
+res = predictor(img)
+```
+
+HF Hub search: [here](https://huggingface.co/models?search=onnxtr).
+
+Collection: [here](https://huggingface.co/collections/Felix92/onnxtr-66bf213a9f88f7346c90e842)
+
+Or push your own models to the hub:
+
+```python
+from onnxtr.models import parseq, push_to_hf_hub, login_to_hub
+from onnxtr.utils.vocabs import VOCABS
+
+# Login to the hub
+login_to_hub()
+
+# Recogniton model
+model = parseq("~/onnxtr-parseq-multilingual-v1.onnx", vocab=VOCABS["multilingual"])
+push_to_hf_hub(
+    model,
+    model_name="onnxtr-parseq-multilingual-v1",
+    task="recognition",  # The task for which the model is intended [detection, recognition, classification]
+    arch="parseq",  # The name of the model architecture
+    override=False  # Set to `True` if you want to override an existing model / repository
+)
+
+# Detection model
+model = linknet_resnet18("~/onnxtr-linknet-resnet18.onnx")
+push_to_hf_hub(
+    model,
+    model_name="onnxtr-linknet-resnet18",
+    task="detection",
+    arch="linknet_resnet18",
+    override=True
+)
+```
+
 ## Models architectures
 
 Credits where it's due: this repository provides ONNX models for the following architectures, converted from the docTR models:

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr.egg-info/SOURCES.txt

@@ -45,6 +45,8 @@ onnxtr/models/detection/postprocessor/__init__.py
 onnxtr/models/detection/postprocessor/base.py
 onnxtr/models/detection/predictor/__init__.py
 onnxtr/models/detection/predictor/base.py
+onnxtr/models/factory/__init__.py
+onnxtr/models/factory/hub.py
 onnxtr/models/predictor/__init__.py
 onnxtr/models/predictor/base.py
 onnxtr/models/predictor/predictor.py

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr.egg-info/requires.txt

@@ -6,6 +6,7 @@ pyclipper<2.0.0,>=1.2.0
 shapely<3.0.0,>=1.6.0
 rapidfuzz<4.0.0,>=3.0.0
 langdetect<2.0.0,>=1.0.9
+huggingface-hub<1.0.0,>=0.23.0
 Pillow>=9.2.0
 defusedxml>=0.7.0
 anyascii>=0.3.2

{onnxtr-0.3.1 → onnxtr-0.4.0}/pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "shapely>=1.6.0,<3.0.0",
     "rapidfuzz>=3.0.0,<4.0.0",
     "langdetect>=1.0.9,<2.0.0",
+    "huggingface-hub>=0.23.0,<1.0.0",
     "Pillow>=9.2.0",
     "defusedxml>=0.7.0",
     "anyascii>=0.3.2",
@@ -126,6 +127,7 @@ module = [
 	"weasyprint.*",
 	"pypdfium2.*",
 	"langdetect.*",
+    "huggingface_hub.*",
     "rapidfuzz.*",
     "anyascii.*",
     "tqdm.*",

{onnxtr-0.3.1 → onnxtr-0.4.0}/setup.py

@@ -9,7 +9,7 @@ from pathlib import Path
 from setuptools import setup
 
 PKG_NAME = "onnxtr"
-VERSION = os.getenv("BUILD_VERSION", "0.3.1a0")
+VERSION = os.getenv("BUILD_VERSION", "0.4.0a0")
 
 
 if __name__ == "__main__":

onnxtr-0.3.1/onnxtr/version.py

@@ -1 +0,0 @@
-__version__ = 'v0.3.1'

{onnxtr-0.3.1 → onnxtr-0.4.0}/onnxtr/models/preprocessor/base.py

@@ -67,11 +67,12 @@ class PreProcessor(NestedObject):
             if x.dtype not in (np.uint8, np.float32):
                 raise TypeError("unsupported data type for numpy.ndarray")
         x = shape_translate(x, "HWC")
+
+        # Resizing
+        x = self.resize(x)
         # Data type & 255 division
         if x.dtype == np.uint8:
             x = x.astype(np.float32) / 255.0
-        # Resizing
-        x = self.resize(x)
 
         return x
 
@@ -95,13 +96,12 @@ class PreProcessor(NestedObject):
                     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])
-
+            # Data type & 255 division
+            if x.dtype == np.uint8:
+                x = x.astype(np.float32) / 255.0
             batches = [x]
 
         elif isinstance(x, list) and all(isinstance(sample, np.ndarray) for sample in x):