python-doctr 0.12.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

doctr/models/modules/layers/tensorflow.py

@@ -1,171 +0,0 @@
-# Copyright (C) 2021-2025, Mindee.
-
-# 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
-
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras import layers
-
-from doctr.utils.repr import NestedObject
-
-__all__ = ["FASTConvLayer"]
-
-
-class FASTConvLayer(layers.Layer, NestedObject):
-    """Convolutional layer used in the TextNet and FAST architectures"""
-
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int | tuple[int, int],
-        stride: int = 1,
-        dilation: int = 1,
-        groups: int = 1,
-        bias: bool = False,
-    ) -> None:
-        super().__init__()
-
-        self.groups = groups
-        self.in_channels = in_channels
-        self.converted_ks = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
-
-        self.hor_conv, self.hor_bn = None, None
-        self.ver_conv, self.ver_bn = None, None
-
-        padding = ((self.converted_ks[0] - 1) * dilation // 2, (self.converted_ks[1] - 1) * dilation // 2)
-
-        self.activation = layers.ReLU()
-        self.conv_pad = layers.ZeroPadding2D(padding=padding)
-
-        self.conv = layers.Conv2D(
-            filters=out_channels,
-            kernel_size=self.converted_ks,
-            strides=stride,
-            dilation_rate=dilation,
-            groups=groups,
-            use_bias=bias,
-        )
-
-        self.bn = layers.BatchNormalization()
-
-        if self.converted_ks[1] != 1:
-            self.ver_pad = layers.ZeroPadding2D(
-                padding=(int(((self.converted_ks[0] - 1) * dilation) / 2), 0),
-            )
-            self.ver_conv = layers.Conv2D(
-                filters=out_channels,
-                kernel_size=(self.converted_ks[0], 1),
-                strides=stride,
-                dilation_rate=dilation,
-                groups=groups,
-                use_bias=bias,
-            )
-            self.ver_bn = layers.BatchNormalization()
-
-        if self.converted_ks[0] != 1:
-            self.hor_pad = layers.ZeroPadding2D(
-                padding=(0, int(((self.converted_ks[1] - 1) * dilation) / 2)),
-            )
-            self.hor_conv = layers.Conv2D(
-                filters=out_channels,
-                kernel_size=(1, self.converted_ks[1]),
-                strides=stride,
-                dilation_rate=dilation,
-                groups=groups,
-                use_bias=bias,
-            )
-            self.hor_bn = layers.BatchNormalization()
-
-        self.rbr_identity = layers.BatchNormalization() if out_channels == in_channels and stride == 1 else None
-
-    def call(self, x: tf.Tensor, **kwargs: Any) -> tf.Tensor:
-        if hasattr(self, "fused_conv"):
-            return self.activation(self.fused_conv(self.conv_pad(x, **kwargs), **kwargs))
-
-        main_outputs = self.bn(self.conv(self.conv_pad(x, **kwargs), **kwargs), **kwargs)
-        vertical_outputs = (
-            self.ver_bn(self.ver_conv(self.ver_pad(x, **kwargs), **kwargs), **kwargs)
-            if self.ver_conv is not None and self.ver_bn is not None
-            else 0
-        )
-        horizontal_outputs = (
-            self.hor_bn(self.hor_conv(self.hor_pad(x, **kwargs), **kwargs), **kwargs)
-            if self.hor_bn is not None and self.hor_conv is not None
-            else 0
-        )
-        id_out = self.rbr_identity(x, **kwargs) if self.rbr_identity is not None else 0
-
-        return self.activation(main_outputs + vertical_outputs + horizontal_outputs + id_out)
-
-    # The following logic is used to reparametrize the layer
-    # Adapted from: https://github.com/mindee/doctr/blob/main/doctr/models/modules/layers/pytorch.py
-    def _identity_to_conv(self, identity: layers.BatchNormalization) -> tuple[tf.Tensor, tf.Tensor] | tuple[int, int]:
-        if identity is None or not hasattr(identity, "moving_mean") or not hasattr(identity, "moving_variance"):
-            return 0, 0
-        if not hasattr(self, "id_tensor"):
-            input_dim = self.in_channels // self.groups
-            kernel_value = np.zeros((1, 1, input_dim, self.in_channels), dtype=np.float32)
-            for i in range(self.in_channels):
-                kernel_value[0, 0, i % input_dim, i] = 1
-            id_tensor = tf.constant(kernel_value, dtype=tf.float32)
-            self.id_tensor = self._pad_to_mxn_tensor(id_tensor)
-        kernel = self.id_tensor
-        std = tf.sqrt(identity.moving_variance + identity.epsilon)
-        t = tf.reshape(identity.gamma / std, (1, 1, 1, -1))
-        return kernel * t, identity.beta - identity.moving_mean * identity.gamma / std
-
-    def _fuse_bn_tensor(self, conv: layers.Conv2D, bn: layers.BatchNormalization) -> tuple[tf.Tensor, tf.Tensor]:
-        kernel = conv.kernel
-        kernel = self._pad_to_mxn_tensor(kernel)
-        std = tf.sqrt(bn.moving_variance + bn.epsilon)
-        t = tf.reshape(bn.gamma / std, (1, 1, 1, -1))
-        return kernel * t, bn.beta - bn.moving_mean * bn.gamma / std
-
-    def _get_equivalent_kernel_bias(self):
-        kernel_mxn, bias_mxn = self._fuse_bn_tensor(self.conv, self.bn)
-        if self.ver_conv is not None:
-            kernel_mx1, bias_mx1 = self._fuse_bn_tensor(self.ver_conv, self.ver_bn)
-        else:
-            kernel_mx1, bias_mx1 = 0, 0
-        if self.hor_conv is not None:
-            kernel_1xn, bias_1xn = self._fuse_bn_tensor(self.hor_conv, self.hor_bn)
-        else:
-            kernel_1xn, bias_1xn = 0, 0
-        kernel_id, bias_id = self._identity_to_conv(self.rbr_identity)
-        kernel_mxn = kernel_mxn + kernel_mx1 + kernel_1xn + kernel_id
-        bias_mxn = bias_mxn + bias_mx1 + bias_1xn + bias_id
-        return kernel_mxn, bias_mxn
-
-    def _pad_to_mxn_tensor(self, kernel: tf.Tensor) -> tf.Tensor:
-        kernel_height, kernel_width = self.converted_ks
-        height, width = kernel.shape[:2]
-        pad_left_right = tf.maximum(0, (kernel_width - width) // 2)
-        pad_top_down = tf.maximum(0, (kernel_height - height) // 2)
-        return tf.pad(kernel, [[pad_top_down, pad_top_down], [pad_left_right, pad_left_right], [0, 0], [0, 0]])
-
-    def reparameterize_layer(self):
-        kernel, bias = self._get_equivalent_kernel_bias()
-        self.fused_conv = layers.Conv2D(
-            filters=self.conv.filters,
-            kernel_size=self.conv.kernel_size,
-            strides=self.conv.strides,
-            padding=self.conv.padding,
-            dilation_rate=self.conv.dilation_rate,
-            groups=self.conv.groups,
-            use_bias=True,
-        )
-        # build layer to initialize weights and biases
-        self.fused_conv.build(input_shape=(None, None, None, kernel.shape[-2]))
-        self.fused_conv.set_weights([kernel.numpy(), bias.numpy()])
-        for para in self.trainable_variables:
-            para._trainable = False
-        for attr in ["conv", "bn", "ver_conv", "ver_bn", "hor_conv", "hor_bn"]:
-            if hasattr(self, attr):
-                delattr(self, attr)
-
-        if hasattr(self, "rbr_identity"):
-            delattr(self, "rbr_identity")

doctr/models/modules/transformer/tensorflow.py

@@ -1,235 +0,0 @@
-# Copyright (C) 2021-2025, Mindee.
-
-# This program is licensed under the Apache License 2.0.
-# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
-
-import math
-from collections.abc import Callable
-from typing import Any
-
-import tensorflow as tf
-from tensorflow.keras import layers
-
-from doctr.utils.repr import NestedObject
-
-__all__ = ["Decoder", "PositionalEncoding", "EncoderBlock", "PositionwiseFeedForward", "MultiHeadAttention"]
-
-
-class PositionalEncoding(layers.Layer, NestedObject):
-    """Compute positional encoding"""
-
-    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000) -> None:
-        super(PositionalEncoding, self).__init__()
-        self.dropout = layers.Dropout(rate=dropout)
-
-        # Compute the positional encodings once in log space.
-        pe = tf.Variable(tf.zeros((max_len, d_model)))
-        position = tf.cast(
-            tf.expand_dims(tf.experimental.numpy.arange(start=0, stop=max_len), axis=1), dtype=tf.float32
-        )
-        div_term = tf.math.exp(
-            tf.cast(tf.experimental.numpy.arange(start=0, stop=d_model, step=2), dtype=tf.float32)
-            * -(math.log(10000.0) / d_model)
-        )
-        pe = pe.numpy()
-        pe[:, 0::2] = tf.math.sin(position * div_term)
-        pe[:, 1::2] = tf.math.cos(position * div_term)
-        self.pe = tf.expand_dims(tf.convert_to_tensor(pe), axis=0)
-
-    def call(
-        self,
-        x: tf.Tensor,
-        **kwargs: Any,
-    ) -> tf.Tensor:
-        """Forward pass
-
-        Args:
-            x: embeddings (batch, max_len, d_model)
-            **kwargs: additional arguments
-
-        Returns:
-            positional embeddings (batch, max_len, d_model)
-        """
-        if x.dtype == tf.float16:  # amp fix: cast to half
-            x = x + tf.cast(self.pe[:, : x.shape[1]], dtype=tf.half)
-        else:
-            x = x + self.pe[:, : x.shape[1]]
-        return self.dropout(x, **kwargs)
-
-
-@tf.function
-def scaled_dot_product_attention(
-    query: tf.Tensor, key: tf.Tensor, value: tf.Tensor, mask: tf.Tensor | None = None
-) -> tuple[tf.Tensor, tf.Tensor]:
-    """Scaled Dot-Product Attention"""
-    scores = tf.matmul(query, tf.transpose(key, perm=[0, 1, 3, 2])) / math.sqrt(query.shape[-1])
-    if mask is not None:
-        # NOTE: to ensure the ONNX compatibility, tf.where works only with bool type condition
-        scores = tf.where(mask == False, float("-inf"), scores)  # noqa: E712
-    p_attn = tf.nn.softmax(scores, axis=-1)
-    return tf.matmul(p_attn, value), p_attn
-
-
-class PositionwiseFeedForward(layers.Layer, NestedObject):
-    """Position-wise Feed-Forward Network"""
-
-    def __init__(
-        self, d_model: int, ffd: int, dropout=0.1, activation_fct: Callable[[Any], Any] = layers.ReLU()
-    ) -> None:
-        super(PositionwiseFeedForward, self).__init__()
-        self.activation_fct = activation_fct
-
-        self.first_linear = layers.Dense(ffd, kernel_initializer=tf.initializers.he_uniform())
-        self.sec_linear = layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform())
-        self.dropout = layers.Dropout(rate=dropout)
-
-    def call(self, x: tf.Tensor, **kwargs: Any) -> tf.Tensor:
-        x = self.first_linear(x, **kwargs)
-        x = self.activation_fct(x)
-        x = self.dropout(x, **kwargs)
-        x = self.sec_linear(x, **kwargs)
-        x = self.dropout(x, **kwargs)
-        return x
-
-
-class MultiHeadAttention(layers.Layer, NestedObject):
-    """Multi-Head Attention"""
-
-    def __init__(self, num_heads: int, d_model: int, dropout: float = 0.1) -> None:
-        super().__init__()
-        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
-
-        self.d_k = d_model // num_heads
-        self.num_heads = num_heads
-
-        self.linear_layers = [layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform()) for _ in range(3)]
-        self.output_linear = layers.Dense(d_model, kernel_initializer=tf.initializers.he_uniform())
-
-    def call(
-        self,
-        query: tf.Tensor,
-        key: tf.Tensor,
-        value: tf.Tensor,
-        mask: tf.Tensor = None,
-        **kwargs: Any,
-    ) -> tf.Tensor:
-        batch_size = query.shape[0]
-
-        # linear projections of Q, K, V
-        query, key, value = [
-            tf.transpose(
-                tf.reshape(linear(x, **kwargs), shape=[batch_size, -1, self.num_heads, self.d_k]), perm=[0, 2, 1, 3]
-            )
-            for linear, x in zip(self.linear_layers, (query, key, value))
-        ]
-
-        # apply attention on all the projected vectors in batch
-        x, attn = scaled_dot_product_attention(query, key, value, mask=mask)
-
-        # Concat attention heads
-        x = tf.transpose(x, perm=[0, 2, 1, 3])
-        x = tf.reshape(x, shape=[batch_size, -1, self.num_heads * self.d_k])
-
-        return self.output_linear(x, **kwargs)
-
-
-class EncoderBlock(layers.Layer, NestedObject):
-    """Transformer Encoder Block"""
-
-    def __init__(
-        self,
-        num_layers: int,
-        num_heads: int,
-        d_model: int,
-        dff: int,  # hidden dimension of the feedforward network
-        dropout: float,
-        activation_fct: Callable[[Any], Any] = layers.ReLU(),
-    ) -> None:
-        super().__init__()
-
-        self.num_layers = num_layers
-
-        self.layer_norm_input = layers.LayerNormalization(epsilon=1e-5)
-        self.layer_norm_attention = layers.LayerNormalization(epsilon=1e-5)
-        self.layer_norm_output = layers.LayerNormalization(epsilon=1e-5)
-        self.dropout = layers.Dropout(rate=dropout)
-
-        self.attention = [MultiHeadAttention(num_heads, d_model, dropout) for _ in range(self.num_layers)]
-        self.position_feed_forward = [
-            PositionwiseFeedForward(d_model, dff, dropout, activation_fct) for _ in range(self.num_layers)
-        ]
-
-    def call(self, x: tf.Tensor, mask: tf.Tensor | None = None, **kwargs: Any) -> tf.Tensor:
-        output = x
-
-        for i in range(self.num_layers):
-            normed_output = self.layer_norm_input(output, **kwargs)
-            output = output + self.dropout(
-                self.attention[i](normed_output, normed_output, normed_output, mask, **kwargs),
-                **kwargs,
-            )
-            normed_output = self.layer_norm_attention(output, **kwargs)
-            output = output + self.dropout(self.position_feed_forward[i](normed_output, **kwargs), **kwargs)
-
-        # (batch_size, seq_len, d_model)
-        return self.layer_norm_output(output, **kwargs)
-
-
-class Decoder(layers.Layer, NestedObject):
-    """Transformer Decoder"""
-
-    def __init__(
-        self,
-        num_layers: int,
-        num_heads: int,
-        d_model: int,
-        vocab_size: int,
-        dropout: float = 0.2,
-        dff: int = 2048,  # hidden dimension of the feedforward network
-        maximum_position_encoding: int = 50,
-    ) -> None:
-        super(Decoder, self).__init__()
-        self.num_layers = num_layers
-        self.d_model = d_model
-
-        self.layer_norm_input = layers.LayerNormalization(epsilon=1e-5)
-        self.layer_norm_masked_attention = layers.LayerNormalization(epsilon=1e-5)
-        self.layer_norm_attention = layers.LayerNormalization(epsilon=1e-5)
-        self.layer_norm_output = layers.LayerNormalization(epsilon=1e-5)
-
-        self.dropout = layers.Dropout(rate=dropout)
-        self.embed = layers.Embedding(vocab_size, d_model)
-        self.positional_encoding = PositionalEncoding(d_model, dropout, maximum_position_encoding)
-
-        self.attention = [MultiHeadAttention(num_heads, d_model, dropout) for _ in range(self.num_layers)]
-        self.source_attention = [MultiHeadAttention(num_heads, d_model, dropout) for _ in range(self.num_layers)]
-        self.position_feed_forward = [PositionwiseFeedForward(d_model, dff, dropout) for _ in range(self.num_layers)]
-
-    def call(
-        self,
-        tgt: tf.Tensor,
-        memory: tf.Tensor,
-        source_mask: tf.Tensor | None = None,
-        target_mask: tf.Tensor | None = None,
-        **kwargs: Any,
-    ) -> tf.Tensor:
-        tgt = self.embed(tgt, **kwargs) * math.sqrt(self.d_model)
-        pos_enc_tgt = self.positional_encoding(tgt, **kwargs)
-        output = pos_enc_tgt
-
-        for i in range(self.num_layers):
-            normed_output = self.layer_norm_input(output, **kwargs)
-            output = output + self.dropout(
-                self.attention[i](normed_output, normed_output, normed_output, target_mask, **kwargs),
-                **kwargs,
-            )
-            normed_output = self.layer_norm_masked_attention(output, **kwargs)
-            output = output + self.dropout(
-                self.source_attention[i](normed_output, memory, memory, source_mask, **kwargs),
-                **kwargs,
-            )
-            normed_output = self.layer_norm_attention(output, **kwargs)
-            output = output + self.dropout(self.position_feed_forward[i](normed_output, **kwargs), **kwargs)
-
-        # (batch_size, seq_len, d_model)
-        return self.layer_norm_output(output, **kwargs)

doctr/models/modules/vision_transformer/tensorflow.py

@@ -1,100 +0,0 @@
-# Copyright (C) 2021-2025, Mindee.
-
-# This program is licensed under the Apache License 2.0.
-# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
-
-import math
-from typing import Any
-
-import tensorflow as tf
-from tensorflow.keras import layers
-
-from doctr.utils.repr import NestedObject
-
-__all__ = ["PatchEmbedding"]
-
-
-class PatchEmbedding(layers.Layer, NestedObject):
-    """Compute 2D patch embeddings with cls token and positional encoding"""
-
-    def __init__(self, input_shape: tuple[int, int, int], embed_dim: int, patch_size: tuple[int, int]) -> None:
-        super().__init__()
-        height, width, _ = input_shape
-        self.patch_size = patch_size
-        self.interpolate = True if patch_size[0] == patch_size[1] else False
-        self.grid_size = tuple(s // p for s, p in zip((height, width), self.patch_size))
-        self.num_patches = self.grid_size[0] * self.grid_size[1]
-
-        self.cls_token = self.add_weight(shape=(1, 1, embed_dim), initializer="zeros", trainable=True, name="cls_token")
-        self.positions = self.add_weight(
-            shape=(1, self.num_patches + 1, embed_dim),
-            initializer="zeros",
-            trainable=True,
-            name="positions",
-        )
-        self.projection = layers.Conv2D(
-            filters=embed_dim,
-            kernel_size=self.patch_size,
-            strides=self.patch_size,
-            padding="valid",
-            data_format="channels_last",
-            use_bias=True,
-            kernel_initializer="glorot_uniform",
-            bias_initializer="zeros",
-            name="projection",
-        )
-
-    def interpolate_pos_encoding(self, embeddings: tf.Tensor, height: int, width: int) -> tf.Tensor:
-        """100 % borrowed from:
-        https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/modeling_tf_vit.py
-
-        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
-        resolution images.
-
-        Source:
-        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py
-        """
-        seq_len, dim = embeddings.shape[1:]
-        num_patches = seq_len - 1
-
-        num_positions = self.positions.shape[1] - 1
-
-        if num_patches == num_positions and height == width:
-            return self.positions
-        class_pos_embed = self.positions[:, :1]
-        patch_pos_embed = self.positions[:, 1:]
-        h0 = height // self.patch_size[0]
-        w0 = width // self.patch_size[1]
-        patch_pos_embed = tf.image.resize(
-            images=tf.reshape(
-                patch_pos_embed, shape=(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
-            ),
-            size=(h0, w0),
-            method="bilinear",
-        )
-
-        shape = patch_pos_embed.shape
-        assert h0 == shape[-3], "height of interpolated patch embedding doesn't match"
-        assert w0 == shape[-2], "width of interpolated patch embedding doesn't match"
-
-        patch_pos_embed = tf.reshape(tensor=patch_pos_embed, shape=(1, -1, dim))
-        return tf.concat(values=(class_pos_embed, patch_pos_embed), axis=1)
-
-    def call(self, x: tf.Tensor, **kwargs: Any) -> tf.Tensor:
-        B, H, W, C = x.shape
-        assert H % self.patch_size[0] == 0, "Image height must be divisible by patch height"
-        assert W % self.patch_size[1] == 0, "Image width must be divisible by patch width"
-        # patchify image
-        patches = self.projection(x, **kwargs)  # (batch_size, num_patches, d_model)
-        patches = tf.reshape(patches, (B, self.num_patches, -1))  # (batch_size, num_patches, d_model)
-
-        cls_tokens = tf.repeat(self.cls_token, B, axis=0)  # (batch_size, 1, d_model)
-        # concate cls_tokens to patches
-        embeddings = tf.concat([cls_tokens, patches], axis=1)  # (batch_size, num_patches + 1, d_model)
-        # add positions to embeddings
-        if self.interpolate:
-            embeddings += self.interpolate_pos_encoding(embeddings, H, W)
-        else:
-            embeddings += self.positions
-
-        return embeddings  # (batch_size, num_patches + 1, d_model)

doctr/models/predictor/tensorflow.py

@@ -1,155 +0,0 @@
-# Copyright (C) 2021-2025, Mindee.
-
-# 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
-
-import numpy as np
-import tensorflow as tf
-
-from doctr.io.elements import Document
-from doctr.models._utils import get_language
-from doctr.models.detection.predictor import DetectionPredictor
-from doctr.models.recognition.predictor import RecognitionPredictor
-from doctr.utils.geometry import detach_scores
-from doctr.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,
-            detect_orientation,
-            **kwargs,
-        )
-        self.detect_orientation = detect_orientation
-        self.detect_language = detect_language
-
-    def __call__(
-        self,
-        pages: list[np.ndarray | tf.Tensor],
-        **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_dict, 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:
-            general_pages_orientations, origin_pages_orientations = self._get_orientations(pages, seg_maps)
-            orientations = [
-                {"value": orientation_page, "confidence": None} for orientation_page in origin_pages_orientations
-            ]
-        else:
-            orientations = None
-            general_pages_orientations = None
-            origin_pages_orientations = None
-        if self.straighten_pages:
-            pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)
-            # update page shapes after straightening
-            origin_page_shapes = [page.shape[:2] for page in pages]
-
-            # forward again to get predictions on straight pages
-            loc_preds_dict = self.det_predictor(pages, **kwargs)
-
-        assert all(len(loc_pred) == 1 for loc_pred in loc_preds_dict), (
-            "Detection Model in ocr_predictor should output only one class"
-        )
-        loc_preds: list[np.ndarray] = [list(loc_pred.values())[0] for loc_pred in loc_preds_dict]
-        # Detach objectness scores from loc_preds
-        loc_preds, objectness_scores = detach_scores(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,
-            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 = []
-        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,
-            objectness_scores,
-            text_preds,
-            origin_page_shapes,
-            crop_orientations,
-            orientations,
-            languages_dict,
-        )
-        return out