onnxtr 0.4.1__py3-none-any.whl → 0.5.1__py3-none-any.whl

onnxtr/io/elements.py

@@ -163,7 +163,7 @@ class Line(Element):
         if geometry is None:
             # Check whether this is a rotated or straight box
             box_resolution_fn = resolve_enclosing_rbbox if len(words[0].geometry) == 4 else resolve_enclosing_bbox
-            geometry = box_resolution_fn([w.geometry for w in words])  # type: ignore[operator]
+            geometry = box_resolution_fn([w.geometry for w in words])  # type: ignore[misc]
 
         super().__init__(words=words)
         self.geometry = geometry
@@ -216,7 +216,7 @@ class Block(Element):
             box_resolution_fn = (
                 resolve_enclosing_rbbox if isinstance(lines[0].geometry, np.ndarray) else resolve_enclosing_bbox
             )
-            geometry = box_resolution_fn(line_boxes + artefact_boxes)  # type: ignore[operator]
+            geometry = box_resolution_fn(line_boxes + artefact_boxes)  # type: ignore
 
         super().__init__(lines=lines, artefacts=artefacts)
         self.geometry = geometry
@@ -294,6 +294,10 @@ class Page(Element):
     def synthesize(self, **kwargs) -> np.ndarray:
         """Synthesize the page from the predictions
 
+        Args:
+        ----
+            **kwargs: keyword arguments passed to the `synthesize_page` method
+
         Returns
         -------
             synthesized page
@@ -442,11 +446,15 @@ class Document(Element):
     def synthesize(self, **kwargs) -> List[np.ndarray]:
         """Synthesize all pages from their predictions
 
+        Args:
+        ----
+            **kwargs: keyword arguments passed to the `Page.synthesize` method
+
         Returns
         -------
             list of synthesized pages
         """
-        return [page.synthesize() for page in self.pages]
+        return [page.synthesize(**kwargs) for page in self.pages]
 
     def export_as_xml(self, **kwargs) -> List[Tuple[bytes, ET.ElementTree]]:
         """Export the document as XML (hOCR-format)

onnxtr/models/builder.py

@@ -266,7 +266,7 @@ class DocumentBuilder(NestedObject):
                 Line([
                     Word(
                         *word_preds[idx],
-                        tuple([tuple(pt) for pt in boxes[idx].tolist()]),  # type: ignore[arg-type]
+                        tuple(tuple(pt) for pt in boxes[idx].tolist()),  # type: ignore[arg-type]
                         float(objectness_scores[idx]),
                         crop_orientations[idx],
                     )

onnxtr/models/classification/models/mobilenet.py

@@ -13,6 +13,7 @@ import numpy as np
 from ...engine import Engine, EngineConfig
 
 __all__ = [
+    "MobileNetV3",
     "mobilenet_v3_small_crop_orientation",
     "mobilenet_v3_small_page_orientation",
 ]

onnxtr/models/classification/predictor/base.py

@@ -3,7 +3,7 @@
 # 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, Union
+from typing import Any, List, Optional, Union
 
 import numpy as np
 from scipy.special import softmax
@@ -29,10 +29,10 @@ class OrientationPredictor(NestedObject):
 
     def __init__(
         self,
-        pre_processor: PreProcessor,
-        model: Any,
+        pre_processor: Optional[PreProcessor],
+        model: Optional[Any],
     ) -> None:
-        self.pre_processor = pre_processor
+        self.pre_processor = pre_processor if isinstance(pre_processor, PreProcessor) else None
         self.model = model
 
     def __call__(
@@ -43,6 +43,10 @@ class OrientationPredictor(NestedObject):
         if any(input.ndim != 3 for input in inputs):
             raise ValueError("incorrect input shape: all inputs are expected to be multi-channel 2D images.")
 
+        if self.model is None or self.pre_processor is None:
+            # predictor is disabled
+            return [[0] * len(inputs), [0] * len(inputs), [1.0] * len(inputs)]
+
         processed_batches = self.pre_processor(inputs)
         predicted_batches = [self.model(batch) for batch in processed_batches]
 

onnxtr/models/classification/zoo.py

@@ -17,16 +17,30 @@ ORIENTATION_ARCHS: List[str] = ["mobilenet_v3_small_crop_orientation", "mobilene
 
 
 def _orientation_predictor(
-    arch: str, load_in_8_bit: bool = False, engine_cfg: Optional[EngineConfig] = None, **kwargs: Any
+    arch: Any,
+    model_type: str,
+    load_in_8_bit: bool = False,
+    engine_cfg: Optional[EngineConfig] = None,
+    disabled: bool = False,
+    **kwargs: Any,
 ) -> OrientationPredictor:
-    if arch not in ORIENTATION_ARCHS:
-        raise ValueError(f"unknown architecture '{arch}'")
+    if disabled:
+        # Case where the orientation predictor is disabled
+        return OrientationPredictor(None, None)
+
+    if isinstance(arch, str):
+        if arch not in ORIENTATION_ARCHS:
+            raise ValueError(f"unknown architecture '{arch}'")
+        # Load directly classifier from backbone
+        _model = classification.__dict__[arch](load_in_8_bit=load_in_8_bit, engine_cfg=engine_cfg)
+    else:
+        if not isinstance(arch, classification.MobileNetV3):
+            raise ValueError(f"unknown architecture: {type(arch)}")
+        _model = arch
 
-    # Load directly classifier from backbone
-    _model = classification.__dict__[arch](load_in_8_bit=load_in_8_bit, engine_cfg=engine_cfg)
     kwargs["mean"] = kwargs.get("mean", _model.cfg["mean"])
     kwargs["std"] = kwargs.get("std", _model.cfg["std"])
-    kwargs["batch_size"] = kwargs.get("batch_size", 512 if "crop" in arch else 2)
+    kwargs["batch_size"] = kwargs.get("batch_size", 512 if model_type == "crop" else 2)
     input_shape = _model.cfg["input_shape"][1:]
     predictor = OrientationPredictor(
         PreProcessor(input_shape, preserve_aspect_ratio=True, symmetric_pad=True, **kwargs),
@@ -60,7 +74,8 @@ def crop_orientation_predictor(
     -------
         OrientationPredictor
     """
-    return _orientation_predictor(arch, load_in_8_bit, engine_cfg, **kwargs)
+    model_type = "crop"
+    return _orientation_predictor(arch, model_type, load_in_8_bit, engine_cfg, **kwargs)
 
 
 def page_orientation_predictor(
@@ -88,4 +103,5 @@ def page_orientation_predictor(
     -------
         OrientationPredictor
     """
-    return _orientation_predictor(arch, load_in_8_bit, engine_cfg, **kwargs)
+    model_type = "page"
+    return _orientation_predictor(arch, model_type, load_in_8_bit, engine_cfg, **kwargs)

onnxtr/models/predictor/base.py

@@ -9,7 +9,7 @@ import numpy as np
 
 from onnxtr.models.builder import DocumentBuilder
 from onnxtr.models.engine import EngineConfig
-from onnxtr.utils.geometry import extract_crops, extract_rcrops, rotate_image
+from onnxtr.utils.geometry import extract_crops, extract_rcrops, remove_image_padding, rotate_image
 
 from .._utils import estimate_orientation, rectify_crops, rectify_loc_preds
 from ..classification import crop_orientation_predictor, page_orientation_predictor
@@ -55,13 +55,19 @@ class _OCRPredictor:
     ) -> None:
         self.assume_straight_pages = assume_straight_pages
         self.straighten_pages = straighten_pages
+        self._page_orientation_disabled = kwargs.pop("disable_page_orientation", False)
+        self._crop_orientation_disabled = kwargs.pop("disable_crop_orientation", False)
         self.crop_orientation_predictor = (
             None
             if assume_straight_pages
-            else crop_orientation_predictor(load_in_8_bit=load_in_8_bit, engine_cfg=clf_engine_cfg)
+            else crop_orientation_predictor(
+                load_in_8_bit=load_in_8_bit, engine_cfg=clf_engine_cfg, disabled=self._crop_orientation_disabled
+            )
         )
         self.page_orientation_predictor = (
-            page_orientation_predictor(load_in_8_bit=load_in_8_bit, engine_cfg=clf_engine_cfg)
+            page_orientation_predictor(
+                load_in_8_bit=load_in_8_bit, engine_cfg=clf_engine_cfg, disabled=self._crop_orientation_disabled
+            )
             if detect_orientation or straighten_pages or not assume_straight_pages
             else None
         )
@@ -112,8 +118,8 @@ class _OCRPredictor:
             ]
         )
         return [
-            # expand if height and width are not equal
-            rotate_image(page, angle, expand=page.shape[0] != page.shape[1])
+            # expand if height and width are not equal, afterwards remove padding
+            remove_image_padding(rotate_image(page, angle, expand=page.shape[0] != page.shape[1]))
             for page, angle in zip(pages, origin_pages_orientations)
         ]
 
@@ -123,13 +129,18 @@ class _OCRPredictor:
         loc_preds: List[np.ndarray],
         channels_last: bool,
         assume_straight_pages: bool = False,
+        assume_horizontal: bool = False,
     ) -> List[List[np.ndarray]]:
-        extraction_fn = extract_crops if assume_straight_pages else extract_rcrops
-
-        crops = [
-            extraction_fn(page, _boxes[:, :4], channels_last=channels_last)  # type: ignore[operator]
-            for page, _boxes in zip(pages, loc_preds)
-        ]
+        if assume_straight_pages:
+            crops = [
+                extract_crops(page, _boxes[:, :4], channels_last=channels_last)
+                for page, _boxes in zip(pages, loc_preds)
+            ]
+        else:
+            crops = [
+                extract_rcrops(page, _boxes[:, :4], channels_last=channels_last, assume_horizontal=assume_horizontal)
+                for page, _boxes in zip(pages, loc_preds)
+            ]
         return crops
 
     @staticmethod
@@ -138,8 +149,9 @@ class _OCRPredictor:
         loc_preds: List[np.ndarray],
         channels_last: bool,
         assume_straight_pages: bool = False,
+        assume_horizontal: bool = False,
     ) -> Tuple[List[List[np.ndarray]], List[np.ndarray]]:
-        crops = _OCRPredictor._generate_crops(pages, loc_preds, channels_last, assume_straight_pages)
+        crops = _OCRPredictor._generate_crops(pages, loc_preds, channels_last, assume_straight_pages, assume_horizontal)
 
         # Avoid sending zero-sized crops
         is_kept = [[all(s > 0 for s in crop.shape) for crop in page_crops] for page_crops in crops]

onnxtr/models/predictor/predictor.py

@@ -119,6 +119,7 @@ class OCRPredictor(NestedObject, _OCRPredictor):
             loc_preds,  # type: ignore[arg-type]
             channels_last=True,
             assume_straight_pages=self.assume_straight_pages,
+            assume_horizontal=self._page_orientation_disabled,
         )
         # Rectify crop orientation and get crop orientation predictions
         crop_orientations: Any = []

onnxtr/utils/geometry.py

@@ -391,6 +391,26 @@ def rotate_image(
     return rot_img
 
 
+def remove_image_padding(image: np.ndarray) -> np.ndarray:
+    """Remove black border padding from an image
+
+    Args:
+    ----
+        image: numpy tensor to remove padding from
+
+    Returns:
+    -------
+        Image with padding removed
+    """
+    # Find the bounding box of the non-black region
+    rows = np.any(image, axis=1)
+    cols = np.any(image, axis=0)
+    rmin, rmax = np.where(rows)[0][[0, -1]]
+    cmin, cmax = np.where(cols)[0][[0, -1]]
+
+    return image[rmin : rmax + 1, cmin : cmax + 1]
+
+
 def estimate_page_angle(polys: np.ndarray) -> float:
     """Takes a batch of rotated previously ORIENTED polys (N, 4, 2) (rectified by the classifier) and return the
     estimated angle ccw in degrees
@@ -471,7 +491,7 @@ def extract_crops(img: np.ndarray, boxes: np.ndarray, channels_last: bool = True
 
 
 def extract_rcrops(
-    img: np.ndarray, polys: np.ndarray, dtype=np.float32, channels_last: bool = True
+    img: np.ndarray, polys: np.ndarray, dtype=np.float32, channels_last: bool = True, assume_horizontal: bool = False
 ) -> List[np.ndarray]:
     """Created cropped images from list of rotated bounding boxes
 
@@ -481,6 +501,7 @@ def extract_rcrops(
         polys: bounding boxes of shape (N, 4, 2)
         dtype: target data type of bounding boxes
         channels_last: whether the channel dimensions is the last one instead of the last one
+        assume_horizontal: whether the boxes are assumed to be only horizontally oriented
 
     Returns:
     -------
@@ -498,22 +519,88 @@ def extract_rcrops(
         _boxes[:, :, 0] *= width
         _boxes[:, :, 1] *= height
 
-    src_pts = _boxes[:, :3].astype(np.float32)
-    # Preserve size
-    d1 = np.linalg.norm(src_pts[:, 0] - src_pts[:, 1], axis=-1)
-    d2 = np.linalg.norm(src_pts[:, 1] - src_pts[:, 2], axis=-1)
-    # (N, 3, 2)
-    dst_pts = np.zeros((_boxes.shape[0], 3, 2), dtype=dtype)
-    dst_pts[:, 1, 0] = dst_pts[:, 2, 0] = d1 - 1
-    dst_pts[:, 2, 1] = d2 - 1
-    # Use a warp transformation to extract the crop
-    crops = [
-        cv2.warpAffine(
-            img if channels_last else img.transpose(1, 2, 0),
-            # Transformation matrix
-            cv2.getAffineTransform(src_pts[idx], dst_pts[idx]),
-            (int(d1[idx]), int(d2[idx])),
-        )
-        for idx in range(_boxes.shape[0])
-    ]
+    src_img = img if channels_last else img.transpose(1, 2, 0)
+
+    # Handle only horizontal oriented boxes
+    if assume_horizontal:
+        crops = []
+
+        for box in _boxes:
+            # Calculate the centroid of the quadrilateral
+            centroid = np.mean(box, axis=0)
+
+            # Divide the points into left and right
+            left_points = box[box[:, 0] < centroid[0]]
+            right_points = box[box[:, 0] >= centroid[0]]
+
+            # Sort the left points according to the y-axis
+            left_points = left_points[np.argsort(left_points[:, 1])]
+            top_left_pt = left_points[0]
+            bottom_left_pt = left_points[-1]
+            # Sort the right points according to the y-axis
+            right_points = right_points[np.argsort(right_points[:, 1])]
+            top_right_pt = right_points[0]
+            bottom_right_pt = right_points[-1]
+            box_points = np.array(
+                [top_left_pt, bottom_left_pt, top_right_pt, bottom_right_pt],
+                dtype=dtype,
+            )
+
+            # Get the width and height of the rectangle that will contain the warped quadrilateral
+            width_upper = np.linalg.norm(top_right_pt - top_left_pt)
+            width_lower = np.linalg.norm(bottom_right_pt - bottom_left_pt)
+            height_left = np.linalg.norm(bottom_left_pt - top_left_pt)
+            height_right = np.linalg.norm(bottom_right_pt - top_right_pt)
+
+            # Get the maximum width and height
+            rect_width = max(int(width_upper), int(width_lower))
+            rect_height = max(int(height_left), int(height_right))
+
+            dst_pts = np.array(
+                [
+                    [0, 0],  # top-left
+                    # bottom-left
+                    [0, rect_height - 1],
+                    # top-right
+                    [rect_width - 1, 0],
+                    # bottom-right
+                    [rect_width - 1, rect_height - 1],
+                ],
+                dtype=dtype,
+            )
+
+            # Get the perspective transform matrix using the box points
+            affine_mat = cv2.getPerspectiveTransform(box_points, dst_pts)
+
+            # Perform the perspective warp to get the rectified crop
+            crop = cv2.warpPerspective(
+                src_img,
+                affine_mat,
+                (rect_width, rect_height),
+            )
+
+            # Add the crop to the list of crops
+            crops.append(crop)
+
+    # Handle any oriented boxes
+    else:
+        src_pts = _boxes[:, :3].astype(np.float32)
+        # Preserve size
+        d1 = np.linalg.norm(src_pts[:, 0] - src_pts[:, 1], axis=-1)
+        d2 = np.linalg.norm(src_pts[:, 1] - src_pts[:, 2], axis=-1)
+        # (N, 3, 2)
+        dst_pts = np.zeros((_boxes.shape[0], 3, 2), dtype=dtype)
+        dst_pts[:, 1, 0] = dst_pts[:, 2, 0] = d1 - 1
+        dst_pts[:, 2, 1] = d2 - 1
+        # Use a warp transformation to extract the crop
+        crops = [
+            cv2.warpAffine(
+                src_img,
+                # Transformation matrix
+                cv2.getAffineTransform(src_pts[idx], dst_pts[idx]),
+                (int(d1[idx]), int(d2[idx])),
+            )
+            for idx in range(_boxes.shape[0])
+        ]
+
     return crops  # type: ignore[return-value]

onnxtr/utils/reconstitution.py

@@ -2,6 +2,7 @@
 
 # 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 typing import Any, Dict, Optional
 
 import numpy as np
@@ -13,10 +14,109 @@ from .fonts import get_font
 __all__ = ["synthesize_page"]
 
 
+# Global variable to avoid multiple warnings
+ROTATION_WARNING = False
+
+
+def _warn_rotation(entry: Dict[str, Any]) -> None:  # pragma: no cover
+    global ROTATION_WARNING
+    if not ROTATION_WARNING and len(entry["geometry"]) == 4:
+        logging.warning("Polygons with larger rotations will lead to inaccurate rendering")
+        ROTATION_WARNING = True
+
+
+def _synthesize(
+    response: Image.Image,
+    entry: Dict[str, Any],
+    w: int,
+    h: int,
+    draw_proba: bool = False,
+    font_family: Optional[str] = None,
+    smoothing_factor: float = 0.75,
+    min_font_size: int = 6,
+    max_font_size: int = 50,
+) -> Image.Image:
+    if len(entry["geometry"]) == 2:
+        (xmin, ymin), (xmax, ymax) = entry["geometry"]
+        polygon = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
+    else:
+        polygon = entry["geometry"]
+
+    # Calculate the bounding box of the word
+    x_coords, y_coords = zip(*polygon)
+    xmin, ymin, xmax, ymax = (
+        int(round(w * min(x_coords))),
+        int(round(h * min(y_coords))),
+        int(round(w * max(x_coords))),
+        int(round(h * max(y_coords))),
+    )
+    word_width = xmax - xmin
+    word_height = ymax - ymin
+
+    # If lines are provided instead of words, concatenate the word entries
+    if "words" in entry:
+        word_text = " ".join(word["value"] for word in entry["words"])
+    else:
+        word_text = entry["value"]
+    # Find the optimal font size
+    try:
+        font_size = min(word_height, max_font_size)
+        font = get_font(font_family, font_size)
+        text_width, text_height = font.getbbox(word_text)[2:4]
+
+        while (text_width > word_width or text_height > word_height) and font_size > min_font_size:
+            font_size = max(int(font_size * smoothing_factor), min_font_size)
+            font = get_font(font_family, font_size)
+            text_width, text_height = font.getbbox(word_text)[2:4]
+    except ValueError:  # pragma: no cover
+        font = get_font(font_family, min_font_size)
+
+    # Create a mask for the word
+    mask = Image.new("L", (w, h), 0)
+    ImageDraw.Draw(mask).polygon([(int(round(w * x)), int(round(h * y))) for x, y in polygon], fill=255)
+
+    # Draw the word text
+    d = ImageDraw.Draw(response)
+    try:
+        try:
+            d.text((xmin, ymin), word_text, font=font, fill=(0, 0, 0), anchor="lt")
+        except UnicodeEncodeError:  # pragma: no cover
+            d.text((xmin, ymin), anyascii(word_text), font=font, fill=(0, 0, 0), anchor="lt")
+    # Catch generic exceptions to avoid crashing the whole rendering
+    except Exception:  # pragma: no cover
+        logging.warning(f"Could not render word: {word_text}")
+
+    if draw_proba:
+        confidence = (
+            entry["confidence"]
+            if "confidence" in entry
+            else sum(w["confidence"] for w in entry["words"]) / len(entry["words"])
+        )
+        p = int(255 * confidence)
+        color = (255 - p, 0, p)  # Red to blue gradient based on probability
+        d.rectangle([(xmin, ymin), (xmax, ymax)], outline=color, width=2)
+
+        prob_font = get_font(font_family, 20)
+        prob_text = f"{confidence:.2f}"
+        prob_text_width, prob_text_height = prob_font.getbbox(prob_text)[2:4]
+
+        # Position the probability slightly above the bounding box
+        prob_x_offset = (word_width - prob_text_width) // 2
+        prob_y_offset = ymin - prob_text_height - 2
+        prob_y_offset = max(0, prob_y_offset)
+
+        d.text((xmin + prob_x_offset, prob_y_offset), prob_text, font=prob_font, fill=color, anchor="lt")
+
+    return response
+
+
 def synthesize_page(
     page: Dict[str, Any],
     draw_proba: bool = False,
     font_family: Optional[str] = None,
+    smoothing_factor: float = 0.95,
+    min_font_size: int = 8,
+    max_font_size: int = 50,
 ) -> np.ndarray:
     """Draw a the content of the element page (OCR response) on a blank page.
 
@@ -24,8 +124,10 @@ def synthesize_page(
     ----
         page: exported Page object to represent
         draw_proba: if True, draw words in colors to represent confidence. Blue: p=1, red: p=0
-        font_size: size of the font, default font = 13
         font_family: family of the font
+        smoothing_factor: factor to smooth the font size
+        min_font_size: minimum font size
+        max_font_size: maximum font size
 
     Returns:
     -------
@@ -33,38 +135,39 @@ def synthesize_page(
     """
     # Draw template
     h, w = page["dimensions"]
-    response = 255 * np.ones((h, w, 3), dtype=np.int32)
+    response = Image.new("RGB", (w, h), color=(255, 255, 255))
 
-    # Draw each word
     for block in page["blocks"]:
-        for line in block["lines"]:
-            for word in line["words"]:
-                # Get absolute word geometry
-                (xmin, ymin), (xmax, ymax) = word["geometry"]
-                xmin, xmax = int(round(w * xmin)), int(round(w * xmax))
-                ymin, ymax = int(round(h * ymin)), int(round(h * ymax))
-
-                # White drawing context adapted to font size, 0.75 factor to convert pts --> pix
-                font = get_font(font_family, int(0.75 * (ymax - ymin)))
-                img = Image.new("RGB", (xmax - xmin, ymax - ymin), color=(255, 255, 255))
-                d = ImageDraw.Draw(img)
-                # Draw in black the value of the word
-                try:
-                    d.text((0, 0), word["value"], font=font, fill=(0, 0, 0))
-                except UnicodeEncodeError:
-                    # When character cannot be encoded, use its anyascii version
-                    d.text((0, 0), anyascii(word["value"]), font=font, fill=(0, 0, 0))
-
-                # Colorize if draw_proba
-                if draw_proba:
-                    p = int(255 * word["confidence"])
-                    mask = np.where(np.array(img) == 0, 1, 0)
-                    proba: np.ndarray = np.array([255 - p, 0, p])
-                    color = mask * proba[np.newaxis, np.newaxis, :]
-                    white_mask = 255 * (1 - mask)
-                    img = color + white_mask
-
-                # Write to response page
-                response[ymin:ymax, xmin:xmax, :] = np.array(img)
+        # If lines are provided use these to get better rendering results
+        if len(block["lines"]) > 1:
+            for line in block["lines"]:
+                _warn_rotation(block)  # pragma: no cover
+                response = _synthesize(
+                    response=response,
+                    entry=line,
+                    w=w,
+                    h=h,
+                    draw_proba=draw_proba,
+                    font_family=font_family,
+                    smoothing_factor=smoothing_factor,
+                    min_font_size=min_font_size,
+                    max_font_size=max_font_size,
+                )
+        # Otherwise, draw each word
+        else:
+            for line in block["lines"]:
+                _warn_rotation(block)  # pragma: no cover
+                for word in line["words"]:
+                    response = _synthesize(
+                        response=response,
+                        entry=word,
+                        w=w,
+                        h=h,
+                        draw_proba=draw_proba,
+                        font_family=font_family,
+                        smoothing_factor=smoothing_factor,
+                        min_font_size=min_font_size,
+                        max_font_size=max_font_size,
+                    )
 
-    return response
+    return np.array(response, dtype=np.uint8)

onnxtr/utils/vocabs.py

@@ -25,6 +25,7 @@ VOCABS: Dict[str, str] = {
     "hindi_punctuation": "।,?!:्ॐ॰॥॰",
     "bangla_letters": "অআইঈউঊঋএঐওঔকখগঘঙচছজঝঞটঠডঢণতথদধনপফবভমযরলশষসহ়ঽািীুূৃেৈোৌ্ৎংঃঁ",
     "bangla_digits": "০১২৩৪৫৬৭৮৯",
+    "generic_cyrillic_letters": "абвгдежзийклмнопрстуфхцчшщьюяАБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЬЮЯ",
 }
 
 VOCABS["latin"] = VOCABS["digits"] + VOCABS["ascii_letters"] + VOCABS["punctuation"]
@@ -59,6 +60,9 @@ VOCABS["vietnamese"] = (
 VOCABS["hebrew"] = VOCABS["english"] + "אבגדהוזחטיכלמנסעפצקרשת" + "₪"
 VOCABS["hindi"] = VOCABS["hindi_letters"] + VOCABS["hindi_digits"] + VOCABS["hindi_punctuation"]
 VOCABS["bangla"] = VOCABS["bangla_letters"] + VOCABS["bangla_digits"]
+VOCABS["ukrainian"] = (
+    VOCABS["generic_cyrillic_letters"] + VOCABS["digits"] + VOCABS["punctuation"] + VOCABS["currency"] + "ґіїєҐІЇЄ₴"
+)
 VOCABS["multilingual"] = "".join(
     dict.fromkeys(
         VOCABS["french"]

onnxtr/version.py

@@ -1 +1 @@
-__version__ = 'v0.4.1'
+__version__ = 'v0.5.1'

{onnxtr-0.4.1.dist-info → onnxtr-0.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxtr
-Version: 0.4.1
+Version: 0.5.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
@@ -227,7 +227,6 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy<3.0.0,>=1.16.0
 Requires-Dist: scipy<2.0.0,>=1.4.0
-Requires-Dist: opencv-python<5.0.0,>=4.5.0
 Requires-Dist: pypdfium2<5.0.0,>=4.11.0
 Requires-Dist: pyclipper<2.0.0,>=1.2.0
 Requires-Dist: shapely<3.0.0,>=1.6.0
@@ -240,8 +239,13 @@ Requires-Dist: anyascii>=0.3.2
 Requires-Dist: tqdm>=4.30.0
 Provides-Extra: cpu
 Requires-Dist: onnxruntime>=1.11.0; extra == "cpu"
+Requires-Dist: opencv-python<5.0.0,>=4.5.0; extra == "cpu"
+Provides-Extra: cpu-headless
+Requires-Dist: onnxruntime>=1.11.0; extra == "cpu-headless"
+Requires-Dist: opencv-python-headless<5.0.0,>=4.5.0; extra == "cpu-headless"
 Provides-Extra: dev
 Requires-Dist: onnxruntime>=1.11.0; extra == "dev"
+Requires-Dist: opencv-python<5.0.0,>=4.5.0; extra == "dev"
 Requires-Dist: weasyprint>=55.0; extra == "dev"
 Requires-Dist: matplotlib>=3.1.0; extra == "dev"
 Requires-Dist: mplcursors>=0.3; extra == "dev"
@@ -253,6 +257,10 @@ Requires-Dist: mypy>=0.812; extra == "dev"
 Requires-Dist: pre-commit>=2.17.0; extra == "dev"
 Provides-Extra: gpu
 Requires-Dist: onnxruntime-gpu>=1.11.0; extra == "gpu"
+Requires-Dist: opencv-python<5.0.0,>=4.5.0; extra == "gpu"
+Provides-Extra: gpu-headless
+Requires-Dist: onnxruntime-gpu>=1.11.0; extra == "gpu-headless"
+Requires-Dist: opencv-python-headless<5.0.0,>=4.5.0; extra == "gpu-headless"
 Provides-Extra: html
 Requires-Dist: weasyprint>=55.0; extra == "html"
 Provides-Extra: quality
@@ -276,7 +284,8 @@ Requires-Dist: mplcursors>=0.3; extra == "viz"
 [![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.4.1-blue.svg)](https://pypi.org/project/OnnxTR/)
+[![Pypi](https://img.shields.io/badge/pypi-v0.5.0-blue.svg)](https://pypi.org/project/OnnxTR/)
+[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Felix92/OnnxTR-OCR)
 
 > :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.
 
@@ -309,8 +318,10 @@ For GPU support please take a look at: [ONNX Runtime](https://onnxruntime.ai/get
 
 ```shell
 pip install "onnxtr[cpu]"
+pip install "onnxtr[cpu-headless]"  # same as cpu but with opencv-headless
 # with gpu support
 pip install "onnxtr[gpu]"
+pip install "onnxtr[gpu-headless]"  # same as gpu but with opencv-headless
 # with HTML support
 pip install "onnxtr[html]"
 # with support for visualization
@@ -356,6 +367,9 @@ model = ocr_predictor(
     # Additional parameters - meta information
     detect_orientation=False,  # set to `True` if the orientation of the pages should be detected (default: False)
     detect_language=False, # set to `True` if the language of the pages should be detected (default: False)
+    # Orientation specific parameters in combination with `assume_straight_pages=False` and/or `straighten_pages=True`
+    disable_crop_orientation=False,  # set to `True` if the crop orientation classification should be disabled (default: False)
+    disable_page_orientation=False,  # set to `True` if the general page orientation classification should be disabled (default: False)
     # DocumentBuilder specific parameters
     resolve_lines=True,  # whether words should be automatically grouped into lines (default: True)
     resolve_blocks=False,  # whether lines should be automatically grouped into blocks (default: False)

{onnxtr-0.4.1.dist-info → onnxtr-0.5.1.dist-info}/RECORD

@@ -1,27 +1,27 @@
 onnxtr/__init__.py,sha256=h7Wc2tuHLsaoCk5xNpEFEK-g11A6SJA7nAasA76TQ_Y,100
 onnxtr/file_utils.py,sha256=WjUKalEdR53aoeIY4e-ihy3r7J_C9qFxL40JHGPfutc,1107
 onnxtr/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-onnxtr/version.py,sha256=FMLDI1-41V4TL8AtTJYDQPs1_pJ9uac4WXnI9j16CjY,23
+onnxtr/version.py,sha256=6swtMqpBHD3aLY2AA1x8huT5k183agnIJ6bjR3m048c,23
 onnxtr/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnxtr/contrib/artefacts.py,sha256=tdmfhvfXVRYEH7uj4_hqf2cuUGoTieyNK8bXsD3zHwo,5383
 onnxtr/contrib/base.py,sha256=KyJ8_zDSKEWSFBszgCbLjEeI7SKg4N_iH_ZQNf90SWQ,3288
 onnxtr/io/__init__.py,sha256=kS7tKGFvzxOCWBOun-Y8n9CsziwRKNynjwpZEUUI03M,106
-onnxtr/io/elements.py,sha256=h-IxpFqXrvg-fOhpnOqpGFLdG-lR-xYYIxk3chy_MN8,17769
+onnxtr/io/elements.py,sha256=GX6rhwg_ByAlL8rAsuLgPAeJ7JsN3_V2o_ETkhh_U68,17977
 onnxtr/io/html.py,sha256=Em_7PjZ56SugJ9bjjcWLCMVe5ee6uUMKeZovNxJFAXw,737
 onnxtr/io/image.py,sha256=4tLTh2bGdA0ohh3a6mV6xD0KqNOtIVi5lJ06XSmeyMI,1759
 onnxtr/io/pdf.py,sha256=tD0klmxI-gkMXp56f_ZXWyPHLsUBKa_xlhNTtGV6tpU,1367
 onnxtr/io/reader.py,sha256=BA7DPhW-Gkmce_ZfzrOl4H3pSXVy2JBeQEuY3pWrBFg,2852
 onnxtr/models/__init__.py,sha256=QTfZlqUyv1d7NUCbGIUFM1DLOOXe-cqHZ7uaKkGdXvk,157
 onnxtr/models/_utils.py,sha256=KncsNcoWqbsxFwduce2STuGHLhv63nXEHv7CMuh6wYA,6606
-onnxtr/models/builder.py,sha256=Bzg-XHZc5k16Ti2XeV9hm4POTHofe581Azq1a3d1O6E,14296
+onnxtr/models/builder.py,sha256=dEZPHkyq-qXo4ZPs9CrufkkWYRrbvxU38ujTx333Mmo,14294
 onnxtr/models/engine.py,sha256=w1vzEduzVDHuxOb0JEkhPp2whrK7ViP03KZiNUNbe4I,4837
 onnxtr/models/zoo.py,sha256=Zcx0mOfMwUR2YAMd7ug06RvXeG2T1PzR2twS6y9X19A,5352
 onnxtr/models/classification/__init__.py,sha256=h1bZs55iLJBMATtzS4ntTKwfD6OGXBiiqGv_hEnOFnE,41
-onnxtr/models/classification/zoo.py,sha256=jzZMf7hKqN9omGAPHJR83rVDaaWhPm-Rk55Xn4bGaIs,3436
+onnxtr/models/classification/zoo.py,sha256=45l0gM3rMrGd_CTxA-OwULA7AOHjTx55_HMidorIjdc,3908
 onnxtr/models/classification/models/__init__.py,sha256=rohbM6ZQslfYchi7feZwwh-sX3XXRUhgtEJQeurAytQ,24
-onnxtr/models/classification/models/mobilenet.py,sha256=l6Ch7ZwL4tqoN94YhSmudY6XYl5fIILzzu4T9JUwZKs,4881
+onnxtr/models/classification/models/mobilenet.py,sha256=rgxkTpRUk_QtU2fAA-Qg3u6y0iWA-zFsmMmYVZRAWiw,4900
 onnxtr/models/classification/predictor/__init__.py,sha256=ERmmOxz_9mUkIuccNbzUa5Y6gVLLVDdyc4cCxbCCUbY,20
-onnxtr/models/classification/predictor/base.py,sha256=Xfaj2XlaJuQ2R81XqF5RB0Wcvzd4wh7Z6j1ifn2niFc,2097
+onnxtr/models/classification/predictor/base.py,sha256=J-EckJM88S4V7f50ukMVgK0tCQEbRZTQ7d1p77x8lUg,2357
 onnxtr/models/detection/__init__.py,sha256=h1bZs55iLJBMATtzS4ntTKwfD6OGXBiiqGv_hEnOFnE,41
 onnxtr/models/detection/core.py,sha256=ZmVDHLJ1l4LQ8rFSKc7enXDkGcOWrcQv4H0SJWyLsag,3584
 onnxtr/models/detection/zoo.py,sha256=5kz4l67Xkr4YTDoI2wDTiI6HSaB926zfua0SZU-Kaw8,2735
@@ -38,8 +38,8 @@ onnxtr/models/detection/predictor/base.py,sha256=bt8M6I14tWC9DYjrFrqg-AU5u670_uP
 onnxtr/models/factory/__init__.py,sha256=cKPoH2V2157lLMTR2zsljG3_IQHziodqR-XK_LG0D_I,19
 onnxtr/models/factory/hub.py,sha256=Fk6pX9VJD422rnVgLh37o136T_0YAsQFzY2dQplDfa4,7176
 onnxtr/models/predictor/__init__.py,sha256=XL25XkRkgyK7mldF-CWhg2MMakSdP5vLpDLwL59hphk,25
-onnxtr/models/predictor/base.py,sha256=ZutI4iNUWk5I5wFfzip89JEDl0SmN7W7hCWRmVec38w,8813
-onnxtr/models/predictor/predictor.py,sha256=pfyTu2qidlPOpXyNKkh20cZefWYUZlF3VEmzrsQr2K8,6368
+onnxtr/models/predictor/base.py,sha256=JnaOMy6pQh0ht2bc0koGqA9rBi_F6mHNyPYGcg0T7C8,9471
+onnxtr/models/predictor/predictor.py,sha256=-ZQfPT70HKTfAP4etywSTy_W-pq5uPDHU00toHuWVcI,6431
 onnxtr/models/preprocessor/__init__.py,sha256=ERmmOxz_9mUkIuccNbzUa5Y6gVLLVDdyc4cCxbCCUbY,20
 onnxtr/models/preprocessor/base.py,sha256=8ZCKsB-o9uRaUm0x4x9FYpYxLXpwHyq2nVv_TlRgaMw,3990
 onnxtr/models/recognition/__init__.py,sha256=h1bZs55iLJBMATtzS4ntTKwfD6OGXBiiqGv_hEnOFnE,41
@@ -61,15 +61,15 @@ onnxtr/utils/__init__.py,sha256=pESRJKtcQyjRxiMgZPhtPYeLbCj-YSGyMVRHTbcMONU,94
 onnxtr/utils/common_types.py,sha256=eC_NyIwbo9qVF33LiNPqHKfyabWq9mYEKD9gAloo5UU,601
 onnxtr/utils/data.py,sha256=Dh0mgeHJhyPwmm63J90uDVmIYbrp63hh1_SnYLnpgJI,4354
 onnxtr/utils/fonts.py,sha256=27v0cojgUrVxNF8Krb1FybSoykoxFy1XjG8lHRUuiEY,1353
-onnxtr/utils/geometry.py,sha256=u9ei6WW8Yd29rtwnrDYercAY-tWkOLkzBd5Oi6NNyDI,17774
+onnxtr/utils/geometry.py,sha256=mYsxRYpMm-UtwmXTcbiSfe2j6-50ZSWAohTcfyi7aZU,20929
 onnxtr/utils/multithreading.py,sha256=30T7AylM3rb52ZEI3Pk1pfB0VYraTbc7yO2vNODVVFY,2011
-onnxtr/utils/reconstitution.py,sha256=Hx1_ddLevKLzuxXc19UelPdsGlAwqi4f6vRSYKHDUB4,2617
+onnxtr/utils/reconstitution.py,sha256=DGb2Isxc2At2GTOO93rbzNvnLRG2vTc5cpdzdijnA8w,6162
 onnxtr/utils/repr.py,sha256=kfbjGL6KymGT8spo2UL4FJXZ0XRwa7CO7Y1dTVR8dIk,2129
 onnxtr/utils/visualization.py,sha256=CX09qvDnNIw3BFW5F3jM4R9OcpLWAeZyoDyTAOGRvls,9925
-onnxtr/utils/vocabs.py,sha256=9Ufmjf7OczWb0931NjWTL7owXLYOKn5x0ulaoVeJGn8,3855
-onnxtr-0.4.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-onnxtr-0.4.1.dist-info/METADATA,sha256=ghYrvvMe49613vZHcXr8qw6e8o7aoVD4uocdcKDaRYU,31723
-onnxtr-0.4.1.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
-onnxtr-0.4.1.dist-info/top_level.txt,sha256=r_MSUTpspp4pWEEWvly-s7ZkfCg1KwrK6-kBlXkWKU8,7
-onnxtr-0.4.1.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-onnxtr-0.4.1.dist-info/RECORD,,
+onnxtr/utils/vocabs.py,sha256=KGGsSLjGl9YLbAYcVCloNR5OIwMKMUc4idpn08EqYYY,4160
+onnxtr-0.5.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+onnxtr-0.5.1.dist-info/METADATA,sha256=VvjTLpK86BT_Psv02B0MbC9GbFflt33d93uRSf3q5wE,32873
+onnxtr-0.5.1.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
+onnxtr-0.5.1.dist-info/top_level.txt,sha256=r_MSUTpspp4pWEEWvly-s7ZkfCg1KwrK6-kBlXkWKU8,7
+onnxtr-0.5.1.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+onnxtr-0.5.1.dist-info/RECORD,,

{onnxtr-0.4.1.dist-info → onnxtr-0.5.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.1)
+Generator: setuptools (75.2.0)
 Root-Is-Purelib: true
 Tag: py3-none-any