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

doctr/file_utils.py

@@ -3,93 +3,13 @@
 # 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.
 
-# Adapted from https://github.com/huggingface/transformers/blob/master/src/transformers/file_utils.py
-
 import importlib.metadata
-import importlib.util
 import logging
-import os
-
-CLASS_NAME: str = "words"
 
+__all__ = ["requires_package", "CLASS_NAME"]
 
-__all__ = ["is_tf_available", "is_torch_available", "requires_package", "CLASS_NAME"]
-
+CLASS_NAME: str = "words"
 ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
-ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
-
-USE_TF = os.environ.get("USE_TF", "AUTO").upper()
-USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
-
-
-if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES:
-    _torch_available = importlib.util.find_spec("torch") is not None
-    if _torch_available:
-        try:
-            _torch_version = importlib.metadata.version("torch")
-            logging.info(f"PyTorch version {_torch_version} available.")
-        except importlib.metadata.PackageNotFoundError:  # pragma: no cover
-            _torch_available = False
-else:  # pragma: no cover
-    logging.info("Disabling PyTorch because USE_TF is set")
-    _torch_available = False
-
-# Compatibility fix to make sure tensorflow.keras stays at Keras 2
-if "TF_USE_LEGACY_KERAS" not in os.environ:
-    os.environ["TF_USE_LEGACY_KERAS"] = "1"
-
-elif os.environ["TF_USE_LEGACY_KERAS"] != "1":
-    raise ValueError(
-        "docTR is only compatible with Keras 2, but you have explicitly set `TF_USE_LEGACY_KERAS` to `0`. "
-    )
-
-
-def ensure_keras_v2() -> None:  # pragma: no cover
-    if not os.environ.get("TF_USE_LEGACY_KERAS") == "1":
-        os.environ["TF_USE_LEGACY_KERAS"] = "1"
-
-
-if USE_TF in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TORCH not in ENV_VARS_TRUE_VALUES:
-    _tf_available = importlib.util.find_spec("tensorflow") is not None
-    if _tf_available:
-        candidates = (
-            "tensorflow",
-            "tensorflow-cpu",
-            "tensorflow-gpu",
-            "tf-nightly",
-            "tf-nightly-cpu",
-            "tf-nightly-gpu",
-            "intel-tensorflow",
-            "tensorflow-rocm",
-            "tensorflow-macos",
-        )
-        _tf_version = None
-        # For the metadata, we have to look for both tensorflow and tensorflow-cpu
-        for pkg in candidates:
-            try:
-                _tf_version = importlib.metadata.version(pkg)
-                break
-            except importlib.metadata.PackageNotFoundError:
-                pass
-        _tf_available = _tf_version is not None
-    if _tf_available:
-        if int(_tf_version.split(".")[0]) < 2:  # type: ignore[union-attr]  # pragma: no cover
-            logging.info(f"TensorFlow found but with version {_tf_version}. DocTR requires version 2 minimum.")
-            _tf_available = False
-        else:
-            logging.info(f"TensorFlow version {_tf_version} available.")
-            ensure_keras_v2()
-
-else:  # pragma: no cover
-    logging.info("Disabling Tensorflow because USE_TORCH is set")
-    _tf_available = False
-
-
-if not _torch_available and not _tf_available:  # pragma: no cover
-    raise ModuleNotFoundError(
-        "DocTR requires either TensorFlow or PyTorch to be installed. Please ensure one of them"
-        " is installed and that either USE_TF or USE_TORCH is enabled."
-    )
 
 
 def requires_package(name: str, extra_message: str | None = None) -> None:  # pragma: no cover
@@ -108,13 +28,3 @@ def requires_package(name: str, extra_message: str | None = None) -> None:  # pr
             f"\n\n{extra_message if extra_message is not None else ''} "
             f"\nPlease install it with the following command: pip install {name}\n"
         )
-
-
-def is_torch_available():
-    """Whether PyTorch is installed."""
-    return _torch_available
-
-
-def is_tf_available():
-    """Whether TensorFlow is installed."""
-    return _tf_available

doctr/io/elements.py

@@ -347,7 +347,7 @@ class Page(Element):
         )
         # Create the body
         body = SubElement(page_hocr, "body")
-        SubElement(
+        page_div = SubElement(
             body,
             "div",
             attrib={
@@ -362,7 +362,7 @@ class Page(Element):
                 raise TypeError("XML export is only available for straight bounding boxes for now.")
             (xmin, ymin), (xmax, ymax) = block.geometry
             block_div = SubElement(
-                body,
+                page_div,
                 "div",
                 attrib={
                     "class": "ocr_carea",
@@ -550,7 +550,41 @@ class KIEPage(Element):
                         {int(round(xmax * width))} {int(round(ymax * height))}",
                     },
                 )
-                prediction_div.text = prediction.value
+                # NOTE: ocr_par, ocr_line and ocrx_word are the same because the KIE predictions contain only words
+                # This is a workaround to make it PDF/A compatible
+                par_div = SubElement(
+                    prediction_div,
+                    "p",
+                    attrib={
+                        "class": "ocr_par",
+                        "id": f"{class_name}_par_{prediction_count}",
+                        "title": f"bbox {int(round(xmin * width))} {int(round(ymin * height))} \
+                        {int(round(xmax * width))} {int(round(ymax * height))}",
+                    },
+                )
+                line_span = SubElement(
+                    par_div,
+                    "span",
+                    attrib={
+                        "class": "ocr_line",
+                        "id": f"{class_name}_line_{prediction_count}",
+                        "title": f"bbox {int(round(xmin * width))} {int(round(ymin * height))} \
+                        {int(round(xmax * width))} {int(round(ymax * height))}; \
+                        baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0",
+                    },
+                )
+                word_div = SubElement(
+                    line_span,
+                    "span",
+                    attrib={
+                        "class": "ocrx_word",
+                        "id": f"{class_name}_word_{prediction_count}",
+                        "title": f"bbox {int(round(xmin * width))} {int(round(ymin * height))} \
+                        {int(round(xmax * width))} {int(round(ymax * height))}; \
+                        x_wconf {int(round(prediction.confidence * 100))}",
+                    },
+                )
+                word_div.text = prediction.value
                 prediction_count += 1
 
         return ET.tostring(page_hocr, encoding="utf-8", method="xml"), ET.ElementTree(page_hocr)

doctr/io/image/__init__.py

@@ -1,8 +1,2 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
 from .base import *
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/io/image/pytorch.py

@@ -95,4 +95,4 @@ def tensor_from_numpy(npy_img: np.ndarray, dtype: torch.dtype = torch.float32) -
 
 def get_img_shape(img: torch.Tensor) -> tuple[int, int]:
     """Get the shape of an image"""
-    return img.shape[-2:]
+    return img.shape[-2:]  # type: ignore[return-value]

doctr/models/_utils.py

@@ -63,7 +63,7 @@ def estimate_orientation(
         thresh = img.astype(np.uint8)
 
     page_orientation, orientation_confidence = general_page_orientation or (None, 0.0)
-    if page_orientation and orientation_confidence >= min_confidence:
+    if page_orientation is not None and orientation_confidence >= min_confidence:
         # We rotate the image to the general orientation which improves the detection
         # No expand needed bitmap is already padded
         thresh = rotate_image(thresh, -page_orientation)
@@ -87,7 +87,7 @@ def estimate_orientation(
 
     angles = []
     for contour in contours[:n_ct]:
-        _, (w, h), angle = cv2.minAreaRect(contour)  # type: ignore[assignment]
+        _, (w, h), angle = cv2.minAreaRect(contour)
         if w / h > ratio_threshold_for_lines:  # select only contours with ratio like lines
             angles.append(angle)
         elif w / h < 1 / ratio_threshold_for_lines:  # if lines are vertical, substract 90 degree
@@ -100,7 +100,7 @@ def estimate_orientation(
         estimated_angle = -round(median) if abs(median) != 0 else 0
 
     # combine with the general orientation and the estimated angle
-    if page_orientation and orientation_confidence >= min_confidence:
+    if page_orientation is not None and orientation_confidence >= min_confidence:
         # special case where the estimated angle is mostly wrong:
         # case 1: - and + swapped
         # case 2: estimated angle is completely wrong
@@ -184,7 +184,7 @@ def invert_data_structure(
         dictionary of list when x is a list of dictionaries or a list of dictionaries when x is dictionary of lists
     """
     if isinstance(x, dict):
-        assert len({len(v) for v in x.values()}) == 1, "All the lists in the dictionnary should have the same length."
+        assert len({len(v) for v in x.values()}) == 1, "All the lists in the dictionary should have the same length."
         return [dict(zip(x, t)) for t in zip(*x.values())]
     elif isinstance(x, list):
         return {k: [dic[k] for dic in x] for k in x[0]}

doctr/models/classification/__init__.py

@@ -4,4 +4,5 @@ from .vgg import *
 from .magc_resnet import *
 from .vit import *
 from .textnet import *
+from .vip import *
 from .zoo import *

doctr/models/classification/magc_resnet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/magc_resnet/pytorch.py

@@ -14,8 +14,7 @@ from torch import nn
 
 from doctr.datasets import VOCABS
 
-from ...utils.pytorch import load_pretrained_params
-from ..resnet.pytorch import ResNet
+from ..resnet import ResNet
 
 __all__ = ["magc_resnet31"]
 
@@ -73,7 +72,7 @@ class MAGC(nn.Module):
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         batch, _, height, width = inputs.size()
         # (N * headers, C / headers, H , W)
-        x = inputs.view(batch * self.headers, self.single_header_inplanes, height, width)
+        x = inputs.contiguous().view(batch * self.headers, self.single_header_inplanes, height, width)
         shortcut = x
         # (N * headers, C / headers, H * W)
         shortcut = shortcut.view(batch * self.headers, self.single_header_inplanes, height * width)
@@ -136,7 +135,7 @@ def _magc_resnet(
         # The number of classes is not the same as the number of classes in the pretrained model =>
         # remove the last layer weights
         _ignore_keys = ignore_keys if kwargs["num_classes"] != len(default_cfgs[arch]["classes"]) else None
-        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
 
     return model
 

doctr/models/classification/mobilenet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/mobilenet/pytorch.py

@@ -5,6 +5,7 @@
 
 # Greatly inspired by https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenetv3.py
 
+import types
 from copy import deepcopy
 from typing import Any
 
@@ -99,12 +100,25 @@ def _mobilenet_v3(
                 m = getattr(m, child)
             m.stride = (2, 1)
 
+    # monkeypatch the model to allow for loading pretrained parameters
+    def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None:  # noqa: D417
+        """Load pretrained parameters onto the model
+
+        Args:
+            path_or_url: the path or URL to the model parameters (checkpoint)
+            **kwargs: additional arguments to be passed to `doctr.models.utils.load_pretrained_params`
+        """
+        load_pretrained_params(self, path_or_url, **kwargs)
+
+    # Bind method to the instance
+    model.from_pretrained = types.MethodType(from_pretrained, model)
+
     # Load pretrained parameters
     if pretrained:
         # The number of classes is not the same as the number of classes in the pretrained model =>
         # remove the last layer weights
         _ignore_keys = ignore_keys if kwargs["num_classes"] != len(default_cfgs[arch]["classes"]) else None
-        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
 
     model.cfg = _cfg
 

doctr/models/classification/predictor/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/models/classification/predictor/pytorch.py

@@ -35,7 +35,7 @@ class OrientationPredictor(nn.Module):
     @torch.inference_mode()
     def forward(
         self,
-        inputs: list[np.ndarray | torch.Tensor],
+        inputs: list[np.ndarray],
     ) -> list[list[int] | list[float]]:
         # Dimension check
         if any(input.ndim != 3 for input in inputs):
@@ -50,7 +50,7 @@ class OrientationPredictor(nn.Module):
         self.model, processed_batches = set_device_and_dtype(
             self.model, processed_batches, _params.device, _params.dtype
         )
-        predicted_batches = [self.model(batch) for batch in processed_batches]  # type: ignore[misc]
+        predicted_batches = [self.model(batch) for batch in processed_batches]
         # confidence
         probs = [
             torch.max(torch.softmax(batch, dim=1), dim=1).values.cpu().detach().numpy() for batch in predicted_batches

doctr/models/classification/resnet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *  # type: ignore[assignment]
+from .pytorch import *

doctr/models/classification/resnet/pytorch.py

@@ -4,6 +4,7 @@
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
 
 
+import types
 from collections.abc import Callable
 from copy import deepcopy
 from typing import Any
@@ -152,6 +153,15 @@ class ResNet(nn.Sequential):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
+    def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None:
+        """Load pretrained parameters onto the model
+
+        Args:
+            path_or_url: the path or URL to the model parameters (checkpoint)
+            **kwargs: additional arguments to be passed to `doctr.models.utils.load_pretrained_params`
+        """
+        load_pretrained_params(self, path_or_url, **kwargs)
+
 
 def _resnet(
     arch: str,
@@ -179,7 +189,7 @@ def _resnet(
         # The number of classes is not the same as the number of classes in the pretrained model =>
         # remove the last layer weights
         _ignore_keys = ignore_keys if kwargs["num_classes"] != len(default_cfgs[arch]["classes"]) else None
-        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
 
     return model
 
@@ -201,12 +211,25 @@ def _tv_resnet(
 
     # Build the model
     model = arch_fn(**kwargs, weights=None)
-    # Load pretrained parameters
+
+    # monkeypatch the model to allow for loading pretrained parameters
+    def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None:  # noqa: D417
+        """Load pretrained parameters onto the model
+
+        Args:
+            path_or_url: the path or URL to the model parameters (checkpoint)
+            **kwargs: additional arguments to be passed to `doctr.models.utils.load_pretrained_params`
+        """
+        load_pretrained_params(self, path_or_url, **kwargs)
+
+    # Bind method to the instance
+    model.from_pretrained = types.MethodType(from_pretrained, model)
+
     if pretrained:
         # The number of classes is not the same as the number of classes in the pretrained model =>
         # remove the last layer weights
         _ignore_keys = ignore_keys if kwargs["num_classes"] != len(default_cfgs[arch]["classes"]) else None
-        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
 
     model.cfg = _cfg
 

doctr/models/classification/textnet/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/textnet/pytorch.py

@@ -11,7 +11,7 @@ from torch import nn
 
 from doctr.datasets import VOCABS
 
-from ...modules.layers.pytorch import FASTConvLayer
+from ...modules.layers import FASTConvLayer
 from ...utils import conv_sequence_pt, load_pretrained_params
 
 __all__ = ["textnet_tiny", "textnet_small", "textnet_base"]
@@ -93,6 +93,15 @@ class TextNet(nn.Sequential):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 
+    def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None:
+        """Load pretrained parameters onto the model
+
+        Args:
+            path_or_url: the path or URL to the model parameters (checkpoint)
+            **kwargs: additional arguments to be passed to `doctr.models.utils.load_pretrained_params`
+        """
+        load_pretrained_params(self, path_or_url, **kwargs)
+
 
 def _textnet(
     arch: str,
@@ -115,7 +124,7 @@ def _textnet(
         # The number of classes is not the same as the number of classes in the pretrained model =>
         # remove the last layer weights
         _ignore_keys = ignore_keys if kwargs["num_classes"] != len(default_cfgs[arch]["classes"]) else None
-        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
 
     model.cfg = _cfg
 

doctr/models/classification/vgg/__init__.py

@@ -1,6 +1 @@
-from doctr.file_utils import is_tf_available, is_torch_available
-
-if is_torch_available():
-    from .pytorch import *
-elif is_tf_available():
-    from .tensorflow import *
+from .pytorch import *

doctr/models/classification/vgg/pytorch.py

@@ -3,6 +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.
 
+import types
 from copy import deepcopy
 from typing import Any
 
@@ -53,12 +54,26 @@ def _vgg(
     # Patch average pool & classification head
     model.avgpool = nn.AdaptiveAvgPool2d((1, 1))
     model.classifier = nn.Linear(512, kwargs["num_classes"])
+
+    # monkeypatch the model to allow for loading pretrained parameters
+    def from_pretrained(self, path_or_url: str, **kwargs: Any) -> None:  # noqa: D417
+        """Load pretrained parameters onto the model
+
+        Args:
+            path_or_url: the path or URL to the model parameters (checkpoint)
+            **kwargs: additional arguments to be passed to `doctr.models.utils.load_pretrained_params`
+        """
+        load_pretrained_params(self, path_or_url, **kwargs)
+
+    # Bind method to the instance
+    model.from_pretrained = types.MethodType(from_pretrained, model)
+
     # Load pretrained parameters
     if pretrained:
         # The number of classes is not the same as the number of classes in the pretrained model =>
         # remove the last layer weights
         _ignore_keys = ignore_keys if kwargs["num_classes"] != len(default_cfgs[arch]["classes"]) else None
-        load_pretrained_params(model, default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
+        model.from_pretrained(default_cfgs[arch]["url"], ignore_keys=_ignore_keys)
 
     model.cfg = _cfg
 

doctr/models/classification/vip/__init__.py

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

doctr/models/classification/vip/layers/__init__.py

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