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

doctr/contrib/__init__.py

@@ -0,0 +1 @@
+from .artefacts import ArtefactDetector

doctr/contrib/artefacts.py

@@ -1,9 +1,9 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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, Dict, List, Optional, Tuple
+from typing import Any
 
 import cv2
 import numpy as np
@@ -14,7 +14,7 @@ from .base import _BasePredictor
 
 __all__ = ["ArtefactDetector"]
 
-default_cfgs: Dict[str, Dict[str, Any]] = {
+default_cfgs: dict[str, dict[str, Any]] = {
     "yolov8_artefact": {
         "input_shape": (3, 1024, 1024),
         "labels": ["bar_code", "qr_code", "logo", "photo"],
@@ -34,7 +34,6 @@ class ArtefactDetector(_BasePredictor):
     >>> results = detector(doc)
 
     Args:
-    ----
         arch: the architecture to use
         batch_size: the batch size to use
         model_path: the path to the model to use
@@ -50,9 +49,9 @@ class ArtefactDetector(_BasePredictor):
         self,
         arch: str = "yolov8_artefact",
         batch_size: int = 2,
-        model_path: Optional[str] = None,
-        labels: Optional[List[str]] = None,
-        input_shape: Optional[Tuple[int, int, int]] = None,
+        model_path: str | None = None,
+        labels: list[str] | None = None,
+        input_shape: tuple[int, int, int] | None = None,
         conf_threshold: float = 0.5,
         iou_threshold: float = 0.5,
         **kwargs: Any,
@@ -66,7 +65,7 @@ class ArtefactDetector(_BasePredictor):
     def preprocess(self, img: np.ndarray) -> np.ndarray:
         return np.transpose(cv2.resize(img, (self.input_shape[2], self.input_shape[1])), (2, 0, 1)) / np.array(255.0)
 
-    def postprocess(self, output: List[np.ndarray], input_images: List[List[np.ndarray]]) -> List[List[Dict[str, Any]]]:
+    def postprocess(self, output: list[np.ndarray], input_images: list[list[np.ndarray]]) -> list[list[dict[str, Any]]]:
         results = []
 
         for batch in zip(output, input_images):
@@ -109,7 +108,6 @@ class ArtefactDetector(_BasePredictor):
         Display the results
 
         Args:
-        ----
             **kwargs: additional keyword arguments to be passed to `plt.show`
         """
         requires_package("matplotlib", "`.show()` requires matplotlib installed")

doctr/contrib/base.py

@@ -1,9 +1,9 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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, List, Optional
+from typing import Any
 
 import numpy as np
 
@@ -16,32 +16,29 @@ class _BasePredictor:
     Base class for all predictors
 
     Args:
-    ----
         batch_size: the batch size to use
         url: the url to use to download a model if needed
         model_path: the path to the model to use
         **kwargs: additional arguments to be passed to `download_from_url`
     """
 
-    def __init__(self, batch_size: int, url: Optional[str] = None, model_path: Optional[str] = None, **kwargs) -> None:
+    def __init__(self, batch_size: int, url: str | None = None, model_path: str | None = None, **kwargs) -> None:
         self.batch_size = batch_size
         self.session = self._init_model(url, model_path, **kwargs)
 
-        self._inputs: List[np.ndarray] = []
-        self._results: List[Any] = []
+        self._inputs: list[np.ndarray] = []
+        self._results: list[Any] = []
 
-    def _init_model(self, url: Optional[str] = None, model_path: Optional[str] = None, **kwargs: Any) -> Any:
+    def _init_model(self, url: str | None = None, model_path: str | None = None, **kwargs: Any) -> Any:
         """
         Download the model from the given url if needed
 
         Args:
-        ----
             url: the url to use
             model_path: the path to the model to use
             **kwargs: additional arguments to be passed to `download_from_url`
 
         Returns:
-        -------
             Any: the ONNX loaded model
         """
         requires_package("onnxruntime", "`.contrib` module requires `onnxruntime` to be installed.")
@@ -57,40 +54,34 @@ class _BasePredictor:
         Preprocess the input image
 
         Args:
-        ----
             img: the input image to preprocess
 
         Returns:
-        -------
             np.ndarray: the preprocessed image
         """
         raise NotImplementedError
 
-    def postprocess(self, output: List[np.ndarray], input_images: List[List[np.ndarray]]) -> Any:
+    def postprocess(self, output: list[np.ndarray], input_images: list[list[np.ndarray]]) -> Any:
         """
         Postprocess the model output
 
         Args:
-        ----
             output: the model output to postprocess
             input_images: the input images used to generate the output
 
         Returns:
-        -------
             Any: the postprocessed output
         """
         raise NotImplementedError
 
-    def __call__(self, inputs: List[np.ndarray]) -> Any:
+    def __call__(self, inputs: list[np.ndarray]) -> Any:
         """
         Call the model on the given inputs
 
         Args:
-        ----
             inputs: the inputs to use
 
         Returns:
-        -------
             Any: the postprocessed output
         """
         self._inputs = inputs

doctr/datasets/cord.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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.
@@ -6,7 +6,7 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any
 
 import numpy as np
 from tqdm import tqdm
@@ -29,10 +29,10 @@ class CORD(VisionDataset):
     >>> img, target = train_set[0]
 
     Args:
-    ----
         train: whether the subset should be the training one
         use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
         recognition_task: whether the dataset should be used for recognition task
+        detection_task: whether the dataset should be used for detection task
         **kwargs: keyword arguments from `VisionDataset`.
     """
 
@@ -53,6 +53,7 @@ class CORD(VisionDataset):
         train: bool = True,
         use_polygons: bool = False,
         recognition_task: bool = False,
+        detection_task: bool = False,
         **kwargs: Any,
     ) -> None:
         url, sha256, name = self.TRAIN if train else self.TEST
@@ -64,13 +65,20 @@ class CORD(VisionDataset):
             pre_transforms=convert_target_to_relative if not recognition_task else None,
             **kwargs,
         )
+        if recognition_task and detection_task:
+            raise ValueError(
+                "`recognition_task` and `detection_task` cannot be set to True simultaneously. "
+                + "To get the whole dataset with boxes and labels leave both parameters to False."
+            )
 
-        # List images
+        # list images
         tmp_root = os.path.join(self.root, "image")
-        self.data: List[Tuple[Union[str, np.ndarray], Union[str, Dict[str, Any]]]] = []
+        self.data: list[tuple[str | np.ndarray, str | dict[str, Any] | np.ndarray]] = []
         self.train = train
         np_dtype = np.float32
-        for img_path in tqdm(iterable=os.listdir(tmp_root), desc="Unpacking CORD", total=len(os.listdir(tmp_root))):
+        for img_path in tqdm(
+            iterable=os.listdir(tmp_root), desc="Preparing and Loading CORD", total=len(os.listdir(tmp_root))
+        ):
             # File existence check
             if not os.path.exists(os.path.join(tmp_root, img_path)):
                 raise FileNotFoundError(f"unable to locate {os.path.join(tmp_root, img_path)}")
@@ -84,7 +92,7 @@ class CORD(VisionDataset):
                         if len(word["text"]) > 0:
                             x = word["quad"]["x1"], word["quad"]["x2"], word["quad"]["x3"], word["quad"]["x4"]
                             y = word["quad"]["y1"], word["quad"]["y2"], word["quad"]["y3"], word["quad"]["y4"]
-                            box: Union[List[float], np.ndarray]
+                            box: list[float] | np.ndarray
                             if use_polygons:
                                 # (x, y) coordinates of top left, top right, bottom right, bottom left corners
                                 box = np.array(
@@ -109,6 +117,8 @@ class CORD(VisionDataset):
                 )
                 for crop, label in zip(crops, list(text_targets)):
                     self.data.append((crop, label))
+            elif detection_task:
+                self.data.append((img_path, np.asarray(box_targets, dtype=int).clip(min=0)))
             else:
                 self.data.append((
                     img_path,

doctr/datasets/datasets/__init__.py

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

doctr/datasets/datasets/base.py

@@ -1,12 +1,13 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 os
 import shutil
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 
@@ -19,15 +20,15 @@ __all__ = ["_AbstractDataset", "_VisionDataset"]
 
 
 class _AbstractDataset:
-    data: List[Any] = []
-    _pre_transforms: Optional[Callable[[Any, Any], Tuple[Any, Any]]] = None
+    data: list[Any] = []
+    _pre_transforms: Callable[[Any, Any], tuple[Any, Any]] | None = None
 
     def __init__(
         self,
-        root: Union[str, Path],
-        img_transforms: Optional[Callable[[Any], Any]] = None,
-        sample_transforms: Optional[Callable[[Any, Any], Tuple[Any, Any]]] = None,
-        pre_transforms: Optional[Callable[[Any, Any], Tuple[Any, Any]]] = None,
+        root: str | Path,
+        img_transforms: Callable[[Any], Any] | None = None,
+        sample_transforms: Callable[[Any, Any], tuple[Any, Any]] | None = None,
+        pre_transforms: Callable[[Any, Any], tuple[Any, Any]] | None = None,
     ) -> None:
         if not Path(root).is_dir():
             raise ValueError(f"expected a path to a reachable folder: {root}")
@@ -41,10 +42,10 @@ class _AbstractDataset:
     def __len__(self) -> int:
         return len(self.data)
 
-    def _read_sample(self, index: int) -> Tuple[Any, Any]:
+    def _read_sample(self, index: int) -> tuple[Any, Any]:
         raise NotImplementedError
 
-    def __getitem__(self, index: int) -> Tuple[Any, Any]:
+    def __getitem__(self, index: int) -> tuple[Any, Any]:
         # Read image
         img, target = self._read_sample(index)
         # Pre-transforms (format conversion at run-time etc.)
@@ -82,7 +83,6 @@ class _VisionDataset(_AbstractDataset):
     """Implements an abstract dataset
 
     Args:
-    ----
         url: URL of the dataset
         file_name: name of the file once downloaded
         file_hash: expected SHA256 of the file
@@ -96,13 +96,13 @@ class _VisionDataset(_AbstractDataset):
     def __init__(
         self,
         url: str,
-        file_name: Optional[str] = None,
-        file_hash: Optional[str] = None,
+        file_name: str | None = None,
+        file_hash: str | None = None,
         extract_archive: bool = False,
         download: bool = False,
         overwrite: bool = False,
-        cache_dir: Optional[str] = None,
-        cache_subdir: Optional[str] = None,
+        cache_dir: str | None = None,
+        cache_subdir: str | None = None,
         **kwargs: Any,
     ) -> None:
         cache_dir = (
@@ -115,7 +115,7 @@ class _VisionDataset(_AbstractDataset):
 
         file_name = file_name if isinstance(file_name, str) else os.path.basename(url)
         # Download the file if not present
-        archive_path: Union[str, Path] = os.path.join(cache_dir, cache_subdir, file_name)
+        archive_path: str | Path = os.path.join(cache_dir, cache_subdir, file_name)
 
         if not os.path.exists(archive_path) and not download:
             raise ValueError("the dataset needs to be downloaded first with download=True")

doctr/datasets/datasets/pytorch.py

@@ -1,11 +1,11 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 os
 from copy import deepcopy
-from typing import Any, List, Tuple
+from typing import Any
 
 import numpy as np
 import torch
@@ -20,7 +20,7 @@ __all__ = ["AbstractDataset", "VisionDataset"]
 class AbstractDataset(_AbstractDataset):
     """Abstract class for all datasets"""
 
-    def _read_sample(self, index: int) -> Tuple[torch.Tensor, Any]:
+    def _read_sample(self, index: int) -> tuple[torch.Tensor, Any]:
         img_name, target = self.data[index]
 
         # Check target
@@ -29,14 +29,14 @@ class AbstractDataset(_AbstractDataset):
             assert "labels" in target, "Target should contain 'labels' key"
         elif isinstance(target, tuple):
             assert len(target) == 2
-            assert isinstance(target[0], str) or isinstance(
-                target[0], np.ndarray
-            ), "first element of the tuple should be a string or a numpy array"
+            assert isinstance(target[0], str) or isinstance(target[0], np.ndarray), (
+                "first element of the tuple should be a string or a numpy array"
+            )
             assert isinstance(target[1], list), "second element of the tuple should be a list"
         else:
-            assert isinstance(target, str) or isinstance(
-                target, np.ndarray
-            ), "Target should be a string or a numpy array"
+            assert isinstance(target, str) or isinstance(target, np.ndarray), (
+                "Target should be a string or a numpy array"
+            )
 
         # Read image
         img = (
@@ -48,11 +48,11 @@ class AbstractDataset(_AbstractDataset):
         return img, deepcopy(target)
 
     @staticmethod
-    def collate_fn(samples: List[Tuple[torch.Tensor, Any]]) -> Tuple[torch.Tensor, List[Any]]:
+    def collate_fn(samples: list[tuple[torch.Tensor, Any]]) -> tuple[torch.Tensor, list[Any]]:
         images, targets = zip(*samples)
-        images = torch.stack(images, dim=0)  # type: ignore[assignment]
+        images = torch.stack(images, dim=0)
 
-        return images, list(targets)  # type: ignore[return-value]
+        return images, list(targets)
 
 
 class VisionDataset(AbstractDataset, _VisionDataset):  # noqa: D101

doctr/datasets/datasets/tensorflow.py

@@ -1,11 +1,11 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 os
 from copy import deepcopy
-from typing import Any, List, Tuple
+from typing import Any
 
 import numpy as np
 import tensorflow as tf
@@ -20,7 +20,7 @@ __all__ = ["AbstractDataset", "VisionDataset"]
 class AbstractDataset(_AbstractDataset):
     """Abstract class for all datasets"""
 
-    def _read_sample(self, index: int) -> Tuple[tf.Tensor, Any]:
+    def _read_sample(self, index: int) -> tuple[tf.Tensor, Any]:
         img_name, target = self.data[index]
 
         # Check target
@@ -29,14 +29,14 @@ class AbstractDataset(_AbstractDataset):
             assert "labels" in target, "Target should contain 'labels' key"
         elif isinstance(target, tuple):
             assert len(target) == 2
-            assert isinstance(target[0], str) or isinstance(
-                target[0], np.ndarray
-            ), "first element of the tuple should be a string or a numpy array"
+            assert isinstance(target[0], str) or isinstance(target[0], np.ndarray), (
+                "first element of the tuple should be a string or a numpy array"
+            )
             assert isinstance(target[1], list), "second element of the tuple should be a list"
         else:
-            assert isinstance(target, str) or isinstance(
-                target, np.ndarray
-            ), "Target should be a string or a numpy array"
+            assert isinstance(target, str) or isinstance(target, np.ndarray), (
+                "Target should be a string or a numpy array"
+            )
 
         # Read image
         img = (
@@ -48,7 +48,7 @@ class AbstractDataset(_AbstractDataset):
         return img, deepcopy(target)
 
     @staticmethod
-    def collate_fn(samples: List[Tuple[tf.Tensor, Any]]) -> Tuple[tf.Tensor, List[Any]]:
+    def collate_fn(samples: list[tuple[tf.Tensor, Any]]) -> tuple[tf.Tensor, list[Any]]:
         images, targets = zip(*samples)
         images = tf.stack(images, axis=0)
 

doctr/datasets/detection.py

@@ -1,11 +1,11 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 json
 import os
-from typing import Any, Dict, List, Tuple, Type, Union
+from typing import Any
 
 import numpy as np
 
@@ -26,7 +26,6 @@ class DetectionDataset(AbstractDataset):
     >>> img, target = train_set[0]
 
     Args:
-    ----
         img_folder: folder with all the images of the dataset
         label_path: path to the annotations of each image
         use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
@@ -47,13 +46,13 @@ class DetectionDataset(AbstractDataset):
         )
 
         # File existence check
-        self._class_names: List = []
+        self._class_names: list = []
         if not os.path.exists(label_path):
             raise FileNotFoundError(f"unable to locate {label_path}")
         with open(label_path, "rb") as f:
             labels = json.load(f)
 
-        self.data: List[Tuple[str, Tuple[np.ndarray, List[str]]]] = []
+        self.data: list[tuple[str, tuple[np.ndarray, list[str]]]] = []
         np_dtype = np.float32
         for img_name, label in labels.items():
             # File existence check
@@ -65,18 +64,16 @@ class DetectionDataset(AbstractDataset):
             self.data.append((img_name, (np.asarray(geoms, dtype=np_dtype), polygons_classes)))
 
     def format_polygons(
-        self, polygons: Union[List, Dict], use_polygons: bool, np_dtype: Type
-    ) -> Tuple[np.ndarray, List[str]]:
+        self, polygons: list | dict, use_polygons: bool, np_dtype: type
+    ) -> tuple[np.ndarray, list[str]]:
         """Format polygons into an array
 
         Args:
-        ----
             polygons: the bounding boxes
             use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
             np_dtype: dtype of array
 
         Returns:
-        -------
             geoms: bounding boxes as np array
             polygons_classes: list of classes for each bounding box
         """

doctr/datasets/doc_artefacts.py

@@ -1,11 +1,11 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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 json
 import os
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 import numpy as np
 
@@ -26,7 +26,6 @@ class DocArtefacts(VisionDataset):
     >>> img, target = train_set[0]
 
     Args:
-    ----
         train: whether the subset should be the training one
         use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
         **kwargs: keyword arguments from `VisionDataset`.
@@ -51,7 +50,7 @@ class DocArtefacts(VisionDataset):
         tmp_root = os.path.join(self.root, "images")
         with open(os.path.join(self.root, "labels.json"), "rb") as f:
             labels = json.load(f)
-        self.data: List[Tuple[str, Dict[str, Any]]] = []
+        self.data: list[tuple[str, dict[str, Any]]] = []
         img_list = os.listdir(tmp_root)
         if len(labels) != len(img_list):
             raise AssertionError("the number of images and labels do not match")

doctr/datasets/funsd.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2024, Mindee.
+# 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.
@@ -6,7 +6,7 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any
 
 import numpy as np
 from tqdm import tqdm
@@ -29,10 +29,10 @@ class FUNSD(VisionDataset):
     >>> img, target = train_set[0]
 
     Args:
-    ----
         train: whether the subset should be the training one
         use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)
         recognition_task: whether the dataset should be used for recognition task
+        detection_task: whether the dataset should be used for detection task
         **kwargs: keyword arguments from `VisionDataset`.
     """
 
@@ -45,6 +45,7 @@ class FUNSD(VisionDataset):
         train: bool = True,
         use_polygons: bool = False,
         recognition_task: bool = False,
+        detection_task: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -55,16 +56,24 @@ class FUNSD(VisionDataset):
             pre_transforms=convert_target_to_relative if not recognition_task else None,
             **kwargs,
         )
+        if recognition_task and detection_task:
+            raise ValueError(
+                "`recognition_task` and `detection_task` cannot be set to True simultaneously. "
+                + "To get the whole dataset with boxes and labels leave both parameters to False."
+            )
+
         self.train = train
         np_dtype = np.float32
 
         # Use the subset
         subfolder = os.path.join("dataset", "training_data" if train else "testing_data")
 
-        # # List images
+        # # list images
         tmp_root = os.path.join(self.root, subfolder, "images")
-        self.data: List[Tuple[Union[str, np.ndarray], Union[str, Dict[str, Any]]]] = []
-        for img_path in tqdm(iterable=os.listdir(tmp_root), desc="Unpacking FUNSD", total=len(os.listdir(tmp_root))):
+        self.data: list[tuple[str | np.ndarray, str | dict[str, Any] | np.ndarray]] = []
+        for img_path in tqdm(
+            iterable=os.listdir(tmp_root), desc="Preparing and Loading FUNSD", total=len(os.listdir(tmp_root))
+        ):
             # File existence check
             if not os.path.exists(os.path.join(tmp_root, img_path)):
                 raise FileNotFoundError(f"unable to locate {os.path.join(tmp_root, img_path)}")
@@ -100,6 +109,8 @@ class FUNSD(VisionDataset):
                     # filter labels with unknown characters
                     if not any(char in label for char in ["☑", "☐", "\uf703", "\uf702"]):
                         self.data.append((crop, label))
+            elif detection_task:
+                self.data.append((img_path, np.asarray(box_targets, dtype=np_dtype)))
             else:
                 self.data.append((
                     img_path,

doctr/datasets/generator/__init__.py

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