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

doctr/models/detection/linknet/tensorflow.py

@@ -1,369 +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.
-
-# Credits: post-processing adapted from https://github.com/xuannianz/DifferentiableBinarization
-
-from copy import deepcopy
-from typing import Any
-
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras import Model, Sequential, layers, losses
-
-from doctr.file_utils import CLASS_NAME
-from doctr.models.classification import resnet18, resnet34, resnet50
-from doctr.models.utils import (
-    IntermediateLayerGetter,
-    _bf16_to_float32,
-    _build_model,
-    conv_sequence,
-    load_pretrained_params,
-)
-from doctr.utils.repr import NestedObject
-
-from .base import LinkNetPostProcessor, _LinkNet
-
-__all__ = ["LinkNet", "linknet_resnet18", "linknet_resnet34", "linknet_resnet50"]
-
-default_cfgs: dict[str, dict[str, Any]] = {
-    "linknet_resnet18": {
-        "mean": (0.798, 0.785, 0.772),
-        "std": (0.264, 0.2749, 0.287),
-        "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/linknet_resnet18-615a82c5.weights.h5&src=0",
-    },
-    "linknet_resnet34": {
-        "mean": (0.798, 0.785, 0.772),
-        "std": (0.264, 0.2749, 0.287),
-        "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/linknet_resnet34-9d772be5.weights.h5&src=0",
-    },
-    "linknet_resnet50": {
-        "mean": (0.798, 0.785, 0.772),
-        "std": (0.264, 0.2749, 0.287),
-        "input_shape": (1024, 1024, 3),
-        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/linknet_resnet50-6bf6c8b5.weights.h5&src=0",
-    },
-}
-
-
-def decoder_block(in_chan: int, out_chan: int, stride: int, **kwargs: Any) -> Sequential:
-    """Creates a LinkNet decoder block"""
-    return Sequential([
-        *conv_sequence(in_chan // 4, "relu", True, kernel_size=1, **kwargs),
-        layers.Conv2DTranspose(
-            filters=in_chan // 4,
-            kernel_size=3,
-            strides=stride,
-            padding="same",
-            use_bias=False,
-            kernel_initializer="he_normal",
-        ),
-        layers.BatchNormalization(),
-        layers.Activation("relu"),
-        *conv_sequence(out_chan, "relu", True, kernel_size=1),
-    ])
-
-
-class LinkNetFPN(Model, NestedObject):
-    """LinkNet Decoder module"""
-
-    def __init__(
-        self,
-        out_chans: int,
-        in_shapes: list[tuple[int, ...]],
-    ) -> None:
-        super().__init__()
-        self.out_chans = out_chans
-        strides = [2] * (len(in_shapes) - 1) + [1]
-        i_chans = [s[-1] for s in in_shapes[::-1]]
-        o_chans = i_chans[1:] + [out_chans]
-        self.decoders = [
-            decoder_block(in_chan, out_chan, s, input_shape=in_shape)
-            for in_chan, out_chan, s, in_shape in zip(i_chans, o_chans, strides, in_shapes[::-1])
-        ]
-
-    def call(self, x: list[tf.Tensor], **kwargs: Any) -> tf.Tensor:
-        out = 0
-        for decoder, fmap in zip(self.decoders, x[::-1]):
-            out = decoder(out + fmap, **kwargs)
-        return out
-
-    def extra_repr(self) -> str:
-        return f"out_chans={self.out_chans}"
-
-
-class LinkNet(_LinkNet, Model):
-    """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
-    <https://arxiv.org/pdf/1707.03718.pdf>`_.
-
-    Args:
-        feature extractor: the backbone serving as feature extractor
-        fpn_channels: number of channels each extracted feature maps is mapped to
-        bin_thresh: threshold for binarization of the output feature map
-        box_thresh: minimal objectness score to consider a box
-        assume_straight_pages: if True, fit straight bounding boxes only
-        exportable: onnx exportable returns only logits
-        cfg: the configuration dict of the model
-        class_names: list of class names
-    """
-
-    _children_names: list[str] = ["feat_extractor", "fpn", "classifier", "postprocessor"]
-
-    def __init__(
-        self,
-        feat_extractor: IntermediateLayerGetter,
-        fpn_channels: int = 64,
-        bin_thresh: float = 0.1,
-        box_thresh: float = 0.1,
-        assume_straight_pages: bool = True,
-        exportable: bool = False,
-        cfg: dict[str, Any] | None = None,
-        class_names: list[str] = [CLASS_NAME],
-    ) -> None:
-        super().__init__(cfg=cfg)
-
-        self.class_names = class_names
-        num_classes: int = len(self.class_names)
-
-        self.exportable = exportable
-        self.assume_straight_pages = assume_straight_pages
-
-        self.feat_extractor = feat_extractor
-
-        self.fpn = LinkNetFPN(fpn_channels, [_shape[1:] for _shape in self.feat_extractor.output_shape])
-        self.fpn.build(self.feat_extractor.output_shape)
-
-        self.classifier = Sequential([
-            layers.Conv2DTranspose(
-                filters=32,
-                kernel_size=3,
-                strides=2,
-                padding="same",
-                use_bias=False,
-                kernel_initializer="he_normal",
-                input_shape=self.fpn.decoders[-1].output_shape[1:],
-            ),
-            layers.BatchNormalization(),
-            layers.Activation("relu"),
-            *conv_sequence(32, "relu", True, kernel_size=3, strides=1),
-            layers.Conv2DTranspose(
-                filters=num_classes,
-                kernel_size=2,
-                strides=2,
-                padding="same",
-                use_bias=True,
-                kernel_initializer="he_normal",
-            ),
-        ])
-
-        self.postprocessor = LinkNetPostProcessor(
-            assume_straight_pages=assume_straight_pages, bin_thresh=bin_thresh, box_thresh=box_thresh
-        )
-
-    def compute_loss(
-        self,
-        out_map: tf.Tensor,
-        target: list[dict[str, np.ndarray]],
-        gamma: float = 2.0,
-        alpha: float = 0.5,
-        eps: float = 1e-8,
-    ) -> tf.Tensor:
-        """Compute linknet loss, BCE with boosted box edges or focal loss. Focal loss implementation based on
-        <https://github.com/tensorflow/addons/>`_.
-
-        Args:
-            out_map: output feature map of the model of shape N x H x W x 1
-            target: list of dictionary where each dict has a `boxes` and a `flags` entry
-            gamma: modulating factor in the focal loss formula
-            alpha: balancing factor in the focal loss formula
-            eps: epsilon factor in dice loss
-
-        Returns:
-            A loss tensor
-        """
-        seg_target, seg_mask = self.build_target(target, out_map.shape[1:], True)
-        seg_target = tf.convert_to_tensor(seg_target, dtype=out_map.dtype)
-        seg_mask = tf.convert_to_tensor(seg_mask, dtype=tf.bool)
-        seg_mask = tf.cast(seg_mask, tf.float32)
-
-        bce_loss = losses.binary_crossentropy(seg_target[..., None], out_map[..., None], from_logits=True)
-        proba_map = tf.sigmoid(out_map)
-
-        # Focal loss
-        if gamma < 0:
-            raise ValueError("Value of gamma should be greater than or equal to zero.")
-        # Convert logits to prob, compute gamma factor
-        p_t = (seg_target * proba_map) + ((1 - seg_target) * (1 - proba_map))
-        alpha_t = seg_target * alpha + (1 - seg_target) * (1 - alpha)
-        # Unreduced loss
-        focal_loss = alpha_t * (1 - p_t) ** gamma * bce_loss
-        # Class reduced
-        focal_loss = tf.reduce_sum(seg_mask * focal_loss, (0, 1, 2, 3)) / tf.reduce_sum(seg_mask, (0, 1, 2, 3))
-
-        # Compute dice loss for each class
-        dice_map = tf.nn.softmax(out_map, axis=-1) if len(self.class_names) > 1 else proba_map
-        # Class-reduced dice loss
-        inter = tf.reduce_sum(seg_mask * dice_map * seg_target, axis=[0, 1, 2])
-        cardinality = tf.reduce_sum(seg_mask * (dice_map + seg_target), axis=[0, 1, 2])
-        dice_loss = tf.reduce_mean(1 - 2 * inter / (cardinality + eps))
-
-        return focal_loss + dice_loss
-
-    def call(
-        self,
-        x: tf.Tensor,
-        target: list[dict[str, np.ndarray]] | None = None,
-        return_model_output: bool = False,
-        return_preds: bool = False,
-        **kwargs: Any,
-    ) -> dict[str, Any]:
-        feat_maps = self.feat_extractor(x, **kwargs)
-        logits = self.fpn(feat_maps, **kwargs)
-        logits = self.classifier(logits, **kwargs)
-
-        out: dict[str, tf.Tensor] = {}
-        if self.exportable:
-            out["logits"] = logits
-            return out
-
-        if return_model_output or target is None or return_preds:
-            prob_map = _bf16_to_float32(tf.math.sigmoid(logits))
-
-        if return_model_output:
-            out["out_map"] = prob_map
-
-        if target is None or return_preds:
-            # Post-process boxes
-            out["preds"] = [dict(zip(self.class_names, preds)) for preds in self.postprocessor(prob_map.numpy())]
-
-        if target is not None:
-            loss = self.compute_loss(logits, target)
-            out["loss"] = loss
-
-        return out
-
-
-def _linknet(
-    arch: str,
-    pretrained: bool,
-    backbone_fn,
-    fpn_layers: list[str],
-    pretrained_backbone: bool = True,
-    input_shape: tuple[int, int, int] | None = None,
-    **kwargs: Any,
-) -> LinkNet:
-    pretrained_backbone = pretrained_backbone and not pretrained
-
-    # Patch the config
-    _cfg = deepcopy(default_cfgs[arch])
-    _cfg["input_shape"] = input_shape or default_cfgs[arch]["input_shape"]
-    if not kwargs.get("class_names", None):
-        kwargs["class_names"] = _cfg.get("class_names", [CLASS_NAME])
-    else:
-        kwargs["class_names"] = sorted(kwargs["class_names"])
-
-    # Feature extractor
-    feat_extractor = IntermediateLayerGetter(
-        backbone_fn(
-            pretrained=pretrained_backbone,
-            include_top=False,
-            input_shape=_cfg["input_shape"],
-        ),
-        fpn_layers,
-    )
-
-    # Build the model
-    model = LinkNet(feat_extractor, cfg=_cfg, **kwargs)
-    _build_model(model)
-
-    # Load pretrained parameters
-    if pretrained:
-        # The given class_names differs from the pretrained model => skip the mismatching layers for fine tuning
-        load_pretrained_params(
-            model,
-            _cfg["url"],
-            skip_mismatch=kwargs["class_names"] != default_cfgs[arch].get("class_names", [CLASS_NAME]),
-        )
-
-    return model
-
-
-def linknet_resnet18(pretrained: bool = False, **kwargs: Any) -> LinkNet:
-    """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
-    <https://arxiv.org/pdf/1707.03718.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import linknet_resnet18
-    >>> model = linknet_resnet18(pretrained=True)
-    >>> input_tensor = tf.random.uniform(shape=[1, 1024, 1024, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on our text detection dataset
-        **kwargs: keyword arguments of the LinkNet architecture
-
-    Returns:
-        text detection architecture
-    """
-    return _linknet(
-        "linknet_resnet18",
-        pretrained,
-        resnet18,
-        ["resnet_block_1", "resnet_block_3", "resnet_block_5", "resnet_block_7"],
-        **kwargs,
-    )
-
-
-def linknet_resnet34(pretrained: bool = False, **kwargs: Any) -> LinkNet:
-    """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
-    <https://arxiv.org/pdf/1707.03718.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import linknet_resnet34
-    >>> model = linknet_resnet34(pretrained=True)
-    >>> input_tensor = tf.random.uniform(shape=[1, 1024, 1024, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on our text detection dataset
-        **kwargs: keyword arguments of the LinkNet architecture
-
-    Returns:
-        text detection architecture
-    """
-    return _linknet(
-        "linknet_resnet34",
-        pretrained,
-        resnet34,
-        ["resnet_block_2", "resnet_block_6", "resnet_block_12", "resnet_block_15"],
-        **kwargs,
-    )
-
-
-def linknet_resnet50(pretrained: bool = False, **kwargs: Any) -> LinkNet:
-    """LinkNet as described in `"LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation"
-    <https://arxiv.org/pdf/1707.03718.pdf>`_.
-
-    >>> import tensorflow as tf
-    >>> from doctr.models import linknet_resnet50
-    >>> model = linknet_resnet50(pretrained=True)
-    >>> input_tensor = tf.random.uniform(shape=[1, 1024, 1024, 3], maxval=1, dtype=tf.float32)
-    >>> out = model(input_tensor)
-
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on our text detection dataset
-        **kwargs: keyword arguments of the LinkNet architecture
-
-    Returns:
-        text detection architecture
-    """
-    return _linknet(
-        "linknet_resnet50",
-        pretrained,
-        resnet50,
-        ["conv2_block3_out", "conv3_block4_out", "conv4_block6_out", "conv5_block3_out"],
-        **kwargs,
-    )

doctr/models/detection/predictor/tensorflow.py

@@ -1,70 +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 Model
-
-from doctr.models.detection._utils import _remove_padding
-from doctr.models.preprocessor import PreProcessor
-from doctr.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: Model,
-    ) -> None:
-        self.pre_processor = pre_processor
-        self.model = model
-
-    def __call__(
-        self,
-        pages: list[np.ndarray | tf.Tensor],
-        return_maps: bool = False,
-        **kwargs: Any,
-    ) -> list[dict[str, np.ndarray]] | tuple[list[dict[str, np.ndarray]], list[np.ndarray]]:
-        # Extract parameters from the preprocessor
-        preserve_aspect_ratio = self.pre_processor.resize.preserve_aspect_ratio
-        symmetric_pad = self.pre_processor.resize.symmetric_pad
-        assume_straight_pages = self.model.assume_straight_pages
-
-        # 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, training=False, **kwargs)
-            for batch in processed_batches
-        ]
-
-        # Remove padding from loc predictions
-        preds = _remove_padding(
-            pages,
-            [pred for batch in predicted_batches for pred in batch["preds"]],
-            preserve_aspect_ratio=preserve_aspect_ratio,
-            symmetric_pad=symmetric_pad,
-            assume_straight_pages=assume_straight_pages,
-        )
-
-        if return_maps:
-            seg_maps = [pred.numpy() for batch in predicted_batches for pred in batch["out_map"]]
-            return preds, seg_maps
-        return preds

doctr/models/kie_predictor/tensorflow.py

@@ -1,187 +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, invert_data_structure
-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 _KIEPredictor
-
-__all__ = ["KIEPredictor"]
-
-
-class KIEPredictor(NestedObject, _KIEPredictor):
-    """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
-        _KIEPredictor.__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, out_maps = self.det_predictor(pages, return_maps=True, **kwargs)
-
-        # Detect document rotation and rotate pages
-        seg_maps = [
-            np.where(np.expand_dims(np.amax(out_map, axis=-1), axis=-1) > kwargs.get("bin_thresh", 0.3), 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 = self.det_predictor(pages, **kwargs)
-
-        dict_loc_preds: dict[str, list[np.ndarray]] = invert_data_structure(loc_preds)  # type: ignore
-
-        # Detach objectness scores from loc_preds
-        objectness_scores = {}
-        for class_name, det_preds in dict_loc_preds.items():
-            _loc_preds, _scores = detach_scores(det_preds)
-            dict_loc_preds[class_name] = _loc_preds
-            objectness_scores[class_name] = _scores
-
-        # Apply hooks to loc_preds if any
-        for hook in self.hooks:
-            dict_loc_preds = hook(dict_loc_preds)
-
-        # Crop images
-        crops = {}
-        for class_name in dict_loc_preds.keys():
-            crops[class_name], dict_loc_preds[class_name] = self._prepare_crops(
-                pages,
-                dict_loc_preds[class_name],
-                channels_last=True,
-                assume_straight_pages=self.assume_straight_pages,
-                assume_horizontal=self._page_orientation_disabled,
-            )
-
-        # Rectify crop orientation
-        crop_orientations: Any = {}
-        if not self.assume_straight_pages:
-            for class_name in dict_loc_preds.keys():
-                crops[class_name], dict_loc_preds[class_name], word_orientations = self._rectify_crops(
-                    crops[class_name], dict_loc_preds[class_name]
-                )
-                crop_orientations[class_name] = [
-                    {"value": orientation[0], "confidence": orientation[1]} for orientation in word_orientations
-                ]
-
-        # Identify character sequences
-        word_preds = {
-            k: self.reco_predictor([crop for page_crops in crop_value for crop in page_crops], **kwargs)
-            for k, crop_value in crops.items()
-        }
-        if not crop_orientations:
-            crop_orientations = {k: [{"value": 0, "confidence": None} for _ in word_preds[k]] for k in word_preds}
-
-        boxes: dict = {}
-        text_preds: dict = {}
-        word_crop_orientations: dict = {}
-        for class_name in dict_loc_preds.keys():
-            boxes[class_name], text_preds[class_name], word_crop_orientations[class_name] = self._process_predictions(
-                dict_loc_preds[class_name], word_preds[class_name], crop_orientations[class_name]
-            )
-
-        boxes_per_page: list[dict] = invert_data_structure(boxes)  # type: ignore[assignment]
-        objectness_scores_per_page: list[dict] = invert_data_structure(objectness_scores)  # type: ignore[assignment]
-        text_preds_per_page: list[dict] = invert_data_structure(text_preds)  # type: ignore[assignment]
-        crop_orientations_per_page: list[dict] = invert_data_structure(word_crop_orientations)  # type: ignore[assignment]
-
-        if self.detect_language:
-            languages = [get_language(self.get_text(text_pred)) for text_pred in text_preds_per_page]
-            languages_dict = [{"value": lang[0], "confidence": lang[1]} for lang in languages]
-        else:
-            languages_dict = None
-
-        out = self.doc_builder(
-            pages,
-            boxes_per_page,
-            objectness_scores_per_page,
-            text_preds_per_page,
-            origin_page_shapes,  # type: ignore[arg-type]
-            crop_orientations_per_page,
-            orientations,
-            languages_dict,
-        )
-        return out
-
-    @staticmethod
-    def get_text(text_pred: dict) -> str:
-        text = []
-        for value in text_pred.values():
-            text += [item[0] for item in value]
-
-        return " ".join(text)