onnxtr 0.1.0__py3-none-any.whl

onnxtr/models/detection/postprocessor/base.py

@@ -0,0 +1,144 @@
+# 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.
+
+# Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
+
+from typing import List, Union
+
+import cv2
+import numpy as np
+import pyclipper
+from shapely.geometry import Polygon
+
+from ..core import DetectionPostProcessor
+
+__all__ = ["GeneralDetectionPostProcessor"]
+
+
+class GeneralDetectionPostProcessor(DetectionPostProcessor):
+    """Implements a post processor for FAST model.
+
+    Args:
+    ----
+        bin_thresh: threshold used to binzarized p_map at inference time
+        box_thresh: minimal objectness score to consider a box
+        assume_straight_pages: whether the inputs were expected to have horizontal text elements
+    """
+
+    def __init__(
+        self,
+        bin_thresh: float = 0.1,
+        box_thresh: float = 0.1,
+        assume_straight_pages: bool = True,
+    ) -> None:
+        super().__init__(box_thresh, bin_thresh, assume_straight_pages)
+        self.unclip_ratio = 1.5
+
+    def polygon_to_box(
+        self,
+        points: np.ndarray,
+    ) -> np.ndarray:
+        """Expand a polygon (points) by a factor unclip_ratio, and returns a polygon
+
+        Args:
+        ----
+            points: The first parameter.
+
+        Returns:
+        -------
+            a box in absolute coordinates (xmin, ymin, xmax, ymax) or (4, 2) array (quadrangle)
+        """
+        if not self.assume_straight_pages:
+            # Compute the rectangle polygon enclosing the raw polygon
+            rect = cv2.minAreaRect(points)
+            points = cv2.boxPoints(rect)
+            # Add 1 pixel to correct cv2 approx
+            area = (rect[1][0] + 1) * (1 + rect[1][1])
+            length = 2 * (rect[1][0] + rect[1][1]) + 2
+        else:
+            poly = Polygon(points)
+            area = poly.area
+            length = poly.length
+        distance = area * self.unclip_ratio / length  # compute distance to expand polygon
+        offset = pyclipper.PyclipperOffset()
+        offset.AddPath(points, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+        _points = offset.Execute(distance)
+        # Take biggest stack of points
+        idx = 0
+        if len(_points) > 1:
+            max_size = 0
+            for _idx, p in enumerate(_points):
+                if len(p) > max_size:
+                    idx = _idx
+                    max_size = len(p)
+            # We ensure that _points can be correctly casted to a ndarray
+            _points = [_points[idx]]
+        expanded_points: np.ndarray = np.asarray(_points)  # expand polygon
+        if len(expanded_points) < 1:
+            return None  # type: ignore[return-value]
+        return (
+            cv2.boundingRect(expanded_points)  # type: ignore[return-value]
+            if self.assume_straight_pages
+            else np.roll(cv2.boxPoints(cv2.minAreaRect(expanded_points)), -1, axis=0)
+        )
+
+    def bitmap_to_boxes(
+        self,
+        pred: np.ndarray,
+        bitmap: np.ndarray,
+    ) -> np.ndarray:
+        """Compute boxes from a bitmap/pred_map: find connected components then filter boxes
+
+        Args:
+        ----
+            pred: Pred map from differentiable linknet output
+            bitmap: Bitmap map computed from pred (binarized)
+            angle_tol: Comparison tolerance of the angle with the median angle across the page
+            ratio_tol: Under this limit aspect ratio, we cannot resolve the direction of the crop
+
+        Returns:
+        -------
+            np tensor boxes for the bitmap, each box is a 6-element list
+                containing x, y, w, h, alpha, score for the box
+        """
+        height, width = bitmap.shape[:2]
+        boxes: List[Union[np.ndarray, List[float]]] = []
+        # get contours from connected components on the bitmap
+        contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        for contour in contours:
+            # Check whether smallest enclosing bounding box is not too small
+            if np.any(contour[:, 0].max(axis=0) - contour[:, 0].min(axis=0) < 2):
+                continue
+            # Compute objectness
+            if self.assume_straight_pages:
+                x, y, w, h = cv2.boundingRect(contour)
+                points: np.ndarray = np.array([[x, y], [x, y + h], [x + w, y + h], [x + w, y]])
+                score = self.box_score(pred, points, assume_straight_pages=True)
+            else:
+                score = self.box_score(pred, contour, assume_straight_pages=False)
+
+            if score < self.box_thresh:  # remove polygons with a weak objectness
+                continue
+
+            if self.assume_straight_pages:
+                _box = self.polygon_to_box(points)
+            else:
+                _box = self.polygon_to_box(np.squeeze(contour))
+
+            if self.assume_straight_pages:
+                # compute relative polygon to get rid of img shape
+                x, y, w, h = _box
+                xmin, ymin, xmax, ymax = x / width, y / height, (x + w) / width, (y + h) / height
+                boxes.append([xmin, ymin, xmax, ymax, score])
+            else:
+                # compute relative box to get rid of img shape
+                _box[:, 0] /= width
+                _box[:, 1] /= height
+                boxes.append(_box)
+
+        if not self.assume_straight_pages:
+            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 4, 2), dtype=pred.dtype)
+        else:
+            return np.clip(np.asarray(boxes), 0, 1) if len(boxes) > 0 else np.zeros((0, 5), dtype=pred.dtype)

onnxtr/models/detection/predictor/__init__.py

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

onnxtr/models/detection/predictor/base.py

@@ -0,0 +1,54 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from typing import Any, List, Tuple, Union
+
+import numpy as np
+
+from onnxtr.models.preprocessor import PreProcessor
+from onnxtr.utils.repr import NestedObject
+
+__all__ = ["DetectionPredictor"]
+
+
+class DetectionPredictor(NestedObject):
+    """Implements an object able to localize text elements in a document
+
+    Args:
+    ----
+        pre_processor: transform inputs for easier batched model inference
+        model: core detection architecture
+    """
+
+    _children_names: List[str] = ["pre_processor", "model"]
+
+    def __init__(
+        self,
+        pre_processor: PreProcessor,
+        model: Any,
+    ) -> None:
+        self.pre_processor = pre_processor
+        self.model = model
+
+    def __call__(
+        self,
+        pages: List[np.ndarray],
+        return_maps: bool = False,
+        **kwargs: Any,
+    ) -> Union[List[np.ndarray], Tuple[List[np.ndarray], List[np.ndarray]]]:
+        # Dimension check
+        if any(page.ndim != 3 for page in pages):
+            raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.")
+
+        processed_batches = self.pre_processor(pages)
+        predicted_batches = [
+            self.model(batch, return_preds=True, return_model_output=True, **kwargs) for batch in processed_batches
+        ]
+
+        preds = [pred for batch in predicted_batches for pred in batch["preds"]]
+        if return_maps:
+            seg_maps = [pred for batch in predicted_batches for pred in batch["out_map"]]
+            return preds, seg_maps
+        return preds

onnxtr/models/detection/zoo.py

@@ -0,0 +1,73 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from typing import Any
+
+from .. import detection
+from ..preprocessor import PreProcessor
+from .predictor import DetectionPredictor
+
+__all__ = ["detection_predictor"]
+
+ARCHS = [
+    "db_resnet34",
+    "db_resnet50",
+    "db_mobilenet_v3_large",
+    "linknet_resnet18",
+    "linknet_resnet34",
+    "linknet_resnet50",
+    "fast_tiny",
+    "fast_small",
+    "fast_base",
+]
+
+
+def _predictor(arch: Any, assume_straight_pages: bool = True, **kwargs: Any) -> DetectionPredictor:
+    if isinstance(arch, str):
+        if arch not in ARCHS:
+            raise ValueError(f"unknown architecture '{arch}'")
+
+        _model = detection.__dict__[arch](assume_straight_pages=assume_straight_pages)
+    else:
+        if not isinstance(arch, (detection.DBNet, detection.LinkNet, detection.FAST)):
+            raise ValueError(f"unknown architecture: {type(arch)}")
+
+        _model = arch
+        _model.postprocessor.assume_straight_pages = assume_straight_pages
+
+    kwargs["mean"] = kwargs.get("mean", _model.cfg["mean"])
+    kwargs["std"] = kwargs.get("std", _model.cfg["std"])
+    kwargs["batch_size"] = kwargs.get("batch_size", 4)
+    predictor = DetectionPredictor(
+        PreProcessor(_model.cfg["input_shape"][1:], **kwargs),
+        _model,
+    )
+    return predictor
+
+
+def detection_predictor(
+    arch: Any = "fast_base",
+    assume_straight_pages: bool = True,
+    **kwargs: Any,
+) -> DetectionPredictor:
+    """Text detection architecture.
+
+    >>> import numpy as np
+    >>> from onnxtr.models import detection_predictor
+    >>> model = detection_predictor(arch='db_resnet50')
+    >>> input_page = (255 * np.random.rand(600, 800, 3)).astype(np.uint8)
+    >>> out = model([input_page])
+
+    Args:
+    ----
+        arch: name of the architecture or model itself to use (e.g. 'db_resnet50')
+        assume_straight_pages: If True, fit straight boxes to the page
+        **kwargs: optional keyword arguments passed to the architecture
+
+    Returns:
+    -------
+        Detection predictor
+    """
+    return _predictor(arch, assume_straight_pages, **kwargs)

onnxtr/models/engine.py

@@ -0,0 +1,50 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from typing import Any, List, Union
+
+import numpy as np
+import onnxruntime
+
+from onnxtr.utils.data import download_from_url
+from onnxtr.utils.geometry import shape_translate
+
+
+class Engine:
+    """Implements an abstract class for the engine of a model
+
+    Args:
+    ----
+        url: the url to use to download a model if needed
+        providers: list of providers to use for inference
+        **kwargs: additional arguments to be passed to `download_from_url`
+    """
+
+    def __init__(
+        self, url: str, providers: List[str] = ["CPUExecutionProvider", "CUDAExecutionProvider"], **kwargs: Any
+    ) -> None:
+        archive_path = download_from_url(url, cache_subdir="models", **kwargs) if "http" in url else url
+        self.runtime = onnxruntime.InferenceSession(archive_path, providers=providers)
+        self.runtime_inputs = self.runtime.get_inputs()[0]
+        self.tf_exported = int(self.runtime_inputs.shape[-1]) == 3
+        self.fixed_batch_size: Union[int, str] = self.runtime_inputs.shape[
+            0
+        ]  # mostly possible with tensorflow exported models
+        self.output_name = [output.name for output in self.runtime.get_outputs()]
+
+    def run(self, inputs: np.ndarray) -> np.ndarray:
+        if self.tf_exported:
+            inputs = shape_translate(inputs, format="BHWC")  # sanity check
+        else:
+            inputs = shape_translate(inputs, format="BCHW")
+        if isinstance(self.fixed_batch_size, int) and self.fixed_batch_size != 0:  # dynamic batch size is a string
+            inputs = np.broadcast_to(inputs, (self.fixed_batch_size, *inputs.shape))
+            # combine the results
+            logits = np.concatenate(
+                [self.runtime.run(self.output_name, {"input": batch})[0] for batch in inputs], axis=0
+            )
+        else:
+            logits = self.runtime.run(self.output_name, {"input": inputs})[0]
+        return shape_translate(logits, format="BHWC")

onnxtr/models/predictor/__init__.py

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

onnxtr/models/predictor/base.py

@@ -0,0 +1,175 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from onnxtr.models.builder import DocumentBuilder
+from onnxtr.utils.geometry import extract_crops, extract_rcrops
+
+from .._utils import rectify_crops, rectify_loc_preds
+from ..classification import crop_orientation_predictor
+from ..classification.predictor import OrientationPredictor
+from ..detection.zoo import ARCHS as DETECTION_ARCHS
+from ..recognition.zoo import ARCHS as RECOGNITION_ARCHS
+
+__all__ = ["_OCRPredictor"]
+
+
+class _OCRPredictor:
+    """Implements an object able to localize and identify text elements in a set of documents
+
+    Args:
+    ----
+        assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
+            without rotated textual elements.
+        straighten_pages: if True, estimates the page general orientation based on the median line orientation.
+            Then, rotates page before passing it to the deep learning modules. The final predictions will be remapped
+            accordingly. Doing so will improve performances for documents with page-uniform rotations.
+        preserve_aspect_ratio: if True, resize preserving the aspect ratio (with padding)
+        symmetric_pad: if True and preserve_aspect_ratio is True, pas the image symmetrically.
+        **kwargs: keyword args of `DocumentBuilder`
+    """
+
+    crop_orientation_predictor: Optional[OrientationPredictor]
+
+    def __init__(
+        self,
+        assume_straight_pages: bool = True,
+        straighten_pages: bool = False,
+        preserve_aspect_ratio: bool = True,
+        symmetric_pad: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        self.assume_straight_pages = assume_straight_pages
+        self.straighten_pages = straighten_pages
+        self.crop_orientation_predictor = None if assume_straight_pages else crop_orientation_predictor()
+        self.doc_builder = DocumentBuilder(**kwargs)
+        self.preserve_aspect_ratio = preserve_aspect_ratio
+        self.symmetric_pad = symmetric_pad
+        self.hooks: List[Callable] = []
+
+    @staticmethod
+    def _generate_crops(
+        pages: List[np.ndarray],
+        loc_preds: List[np.ndarray],
+        channels_last: bool,
+        assume_straight_pages: 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)
+        ]
+        return crops
+
+    @staticmethod
+    def _prepare_crops(
+        pages: List[np.ndarray],
+        loc_preds: List[np.ndarray],
+        channels_last: bool,
+        assume_straight_pages: bool = False,
+    ) -> Tuple[List[List[np.ndarray]], List[np.ndarray]]:
+        crops = _OCRPredictor._generate_crops(pages, loc_preds, channels_last, assume_straight_pages)
+
+        # 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]
+        crops = [
+            [crop for crop, _kept in zip(page_crops, page_kept) if _kept]
+            for page_crops, page_kept in zip(crops, is_kept)
+        ]
+        loc_preds = [_boxes[_kept] for _boxes, _kept in zip(loc_preds, is_kept)]
+
+        return crops, loc_preds
+
+    def _rectify_crops(
+        self,
+        crops: List[List[np.ndarray]],
+        loc_preds: List[np.ndarray],
+    ) -> Tuple[List[List[np.ndarray]], List[np.ndarray], List[Tuple[int, float]]]:
+        # Work at a page level
+        orientations, classes, probs = zip(*[self.crop_orientation_predictor(page_crops) for page_crops in crops])  # type: ignore[misc]
+        rect_crops = [rectify_crops(page_crops, orientation) for page_crops, orientation in zip(crops, orientations)]
+        rect_loc_preds = [
+            rectify_loc_preds(page_loc_preds, orientation) if len(page_loc_preds) > 0 else page_loc_preds
+            for page_loc_preds, orientation in zip(loc_preds, orientations)
+        ]
+        # Flatten to list of tuples with (value, confidence)
+        crop_orientations = [
+            (orientation, prob)
+            for page_classes, page_probs in zip(classes, probs)
+            for orientation, prob in zip(page_classes, page_probs)
+        ]
+        return rect_crops, rect_loc_preds, crop_orientations  # type: ignore[return-value]
+
+    def _remove_padding(
+        self,
+        pages: List[np.ndarray],
+        loc_preds: List[np.ndarray],
+    ) -> List[np.ndarray]:
+        if self.preserve_aspect_ratio:
+            # Rectify loc_preds to remove padding
+            rectified_preds = []
+            for page, loc_pred in zip(pages, loc_preds):
+                h, w = page.shape[0], page.shape[1]
+                if h > w:
+                    # y unchanged, dilate x coord
+                    if self.symmetric_pad:
+                        if self.assume_straight_pages:
+                            loc_pred[:, [0, 2]] = np.clip((loc_pred[:, [0, 2]] - 0.5) * h / w + 0.5, 0, 1)
+                        else:
+                            loc_pred[:, :, 0] = np.clip((loc_pred[:, :, 0] - 0.5) * h / w + 0.5, 0, 1)
+                    else:
+                        if self.assume_straight_pages:
+                            loc_pred[:, [0, 2]] *= h / w
+                        else:
+                            loc_pred[:, :, 0] *= h / w
+                elif w > h:
+                    # x unchanged, dilate y coord
+                    if self.symmetric_pad:
+                        if self.assume_straight_pages:
+                            loc_pred[:, [1, 3]] = np.clip((loc_pred[:, [1, 3]] - 0.5) * w / h + 0.5, 0, 1)
+                        else:
+                            loc_pred[:, :, 1] = np.clip((loc_pred[:, :, 1] - 0.5) * w / h + 0.5, 0, 1)
+                    else:
+                        if self.assume_straight_pages:
+                            loc_pred[:, [1, 3]] *= w / h
+                        else:
+                            loc_pred[:, :, 1] *= w / h
+                rectified_preds.append(loc_pred)
+            return rectified_preds
+        return loc_preds
+
+    @staticmethod
+    def _process_predictions(
+        loc_preds: List[np.ndarray],
+        word_preds: List[Tuple[str, float]],
+        crop_orientations: List[Dict[str, Any]],
+    ) -> Tuple[List[np.ndarray], List[List[Tuple[str, float]]], List[List[Dict[str, Any]]]]:
+        text_preds = []
+        crop_orientation_preds = []
+        if len(loc_preds) > 0:
+            # Text & crop orientation predictions at page level
+            _idx = 0
+            for page_boxes in loc_preds:
+                text_preds.append(word_preds[_idx : _idx + page_boxes.shape[0]])
+                crop_orientation_preds.append(crop_orientations[_idx : _idx + page_boxes.shape[0]])
+                _idx += page_boxes.shape[0]
+
+        return loc_preds, text_preds, crop_orientation_preds
+
+    def add_hook(self, hook: Callable) -> None:
+        """Add a hook to the predictor
+
+        Args:
+        ----
+            hook: a callable that takes as input the `loc_preds` and returns the modified `loc_preds`
+        """
+        self.hooks.append(hook)
+
+    def list_archs(self) -> Dict[str, List[str]]:
+        return {"detection_archs": DETECTION_ARCHS, "recognition_archs": RECOGNITION_ARCHS}

onnxtr/models/predictor/predictor.py

@@ -0,0 +1,145 @@
+# Copyright (C) 2021-2024, Mindee | Felix Dittrich.
+
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
+
+from typing import Any, List
+
+import numpy as np
+
+from onnxtr.io.elements import Document
+from onnxtr.models._utils import estimate_orientation, get_language
+from onnxtr.models.detection.predictor import DetectionPredictor
+from onnxtr.models.recognition.predictor import RecognitionPredictor
+from onnxtr.utils.geometry import rotate_image
+from onnxtr.utils.repr import NestedObject
+
+from .base import _OCRPredictor
+
+__all__ = ["OCRPredictor"]
+
+
+class OCRPredictor(NestedObject, _OCRPredictor):
+    """Implements an object able to localize and identify text elements in a set of documents
+
+    Args:
+    ----
+        det_predictor: detection module
+        reco_predictor: recognition module
+        assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
+            without rotated textual elements.
+        straighten_pages: if True, estimates the page general orientation based on the median line orientation.
+            Then, rotates page before passing it to the deep learning modules. The final predictions will be remapped
+            accordingly. Doing so will improve performances for documents with page-uniform rotations.
+        detect_orientation: if True, the estimated general page orientation will be added to the predictions for each
+            page. Doing so will slightly deteriorate the overall latency.
+        detect_language: if True, the language prediction will be added to the predictions for each
+            page. Doing so will slightly deteriorate the overall latency.
+        **kwargs: keyword args of `DocumentBuilder`
+    """
+
+    _children_names = ["det_predictor", "reco_predictor", "doc_builder"]
+
+    def __init__(
+        self,
+        det_predictor: DetectionPredictor,
+        reco_predictor: RecognitionPredictor,
+        assume_straight_pages: bool = True,
+        straighten_pages: bool = False,
+        preserve_aspect_ratio: bool = True,
+        symmetric_pad: bool = True,
+        detect_orientation: bool = False,
+        detect_language: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        self.det_predictor = det_predictor
+        self.reco_predictor = reco_predictor
+        _OCRPredictor.__init__(
+            self, assume_straight_pages, straighten_pages, preserve_aspect_ratio, symmetric_pad, **kwargs
+        )
+        self.detect_orientation = detect_orientation
+        self.detect_language = detect_language
+
+    def __call__(
+        self,
+        pages: List[np.ndarray],
+        **kwargs: Any,
+    ) -> Document:
+        # Dimension check
+        if any(page.ndim != 3 for page in pages):
+            raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.")
+
+        origin_page_shapes = [page.shape[:2] for page in pages]
+
+        # Localize text elements
+        loc_preds, out_maps = self.det_predictor(pages, return_maps=True, **kwargs)
+
+        # Detect document rotation and rotate pages
+        seg_maps = [
+            np.where(out_map > getattr(self.det_predictor.model.postprocessor, "bin_thresh"), 255, 0).astype(np.uint8)
+            for out_map in out_maps
+        ]
+        if self.detect_orientation:
+            origin_page_orientations = [estimate_orientation(seq_map) for seq_map in seg_maps]
+            orientations = [
+                {"value": orientation_page, "confidence": None} for orientation_page in origin_page_orientations
+            ]
+        else:
+            orientations = None
+        if self.straighten_pages:
+            origin_page_orientations = (
+                origin_page_orientations
+                if self.detect_orientation
+                else [estimate_orientation(seq_map) for seq_map in seg_maps]
+            )
+            pages = [rotate_image(page, -angle, expand=False) for page, angle in zip(pages, origin_page_orientations)]
+            # forward again to get predictions on straight pages
+            loc_preds = self.det_predictor(pages, **kwargs)  # type: ignore[assignment]
+
+        loc_preds = [loc_pred[0] for loc_pred in loc_preds]
+
+        # Rectify crops if aspect ratio
+        loc_preds = self._remove_padding(pages, loc_preds)
+
+        # Apply hooks to loc_preds if any
+        for hook in self.hooks:
+            loc_preds = hook(loc_preds)
+
+        # Crop images
+        crops, loc_preds = self._prepare_crops(
+            pages,
+            loc_preds,  # type: ignore[arg-type]
+            channels_last=True,
+            assume_straight_pages=self.assume_straight_pages,
+        )
+        # Rectify crop orientation and get crop orientation predictions
+        crop_orientations: Any = []
+        if not self.assume_straight_pages:
+            crops, loc_preds, _crop_orientations = self._rectify_crops(crops, loc_preds)
+            crop_orientations = [
+                {"value": orientation[0], "confidence": orientation[1]} for orientation in _crop_orientations
+            ]
+
+        # Identify character sequences
+        word_preds = self.reco_predictor([crop for page_crops in crops for crop in page_crops], **kwargs)
+        if not crop_orientations:
+            crop_orientations = [{"value": 0, "confidence": None} for _ in word_preds]
+
+        boxes, text_preds, crop_orientations = self._process_predictions(loc_preds, word_preds, crop_orientations)
+
+        if self.detect_language:
+            languages = [get_language(" ".join([item[0] for item in text_pred])) for text_pred in text_preds]
+            languages_dict = [{"value": lang[0], "confidence": lang[1]} for lang in languages]
+        else:
+            languages_dict = None
+
+        out = self.doc_builder(
+            pages,
+            boxes,
+            text_preds,
+            origin_page_shapes,  # type: ignore[arg-type]
+            crop_orientations,
+            orientations,
+            languages_dict,
+        )
+        return out

onnxtr/models/preprocessor/__init__.py

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