python-doctr 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

doctr/datasets/__init__.py

@@ -13,12 +13,14 @@ from .imgur5k import *
 from .mjsynth import *
 from .ocr import *
 from .recognition import *
+from .orientation import *
 from .sroie import *
 from .svhn import *
 from .svt import *
 from .synthtext import *
 from .utils import *
 from .vocabs import *
+from .wildreceipt import *
 
 if is_tf_available():
     from .loader import *

doctr/datasets/cord.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -29,6 +29,7 @@ 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
@@ -109,9 +110,10 @@ class CORD(VisionDataset):
                 for crop, label in zip(crops, list(text_targets)):
                     self.data.append((crop, label))
             else:
-                self.data.append(
-                    (img_path, dict(boxes=np.asarray(box_targets, dtype=int).clip(min=0), labels=list(text_targets)))
-                )
+                self.data.append((
+                    img_path,
+                    dict(boxes=np.asarray(box_targets, dtype=int).clip(min=0), labels=list(text_targets)),
+                ))
 
         self.root = tmp_root
 

doctr/datasets/datasets/base.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -59,7 +59,7 @@ class _AbstractDataset:
             # Conditions to assess it is detection model with multiple classes and avoid confusion with other tasks.
             if (
                 isinstance(target, dict)
-                and all([isinstance(item, np.ndarray) for item in target.values()])
+                and all(isinstance(item, np.ndarray) for item in target.values())
                 and set(target.keys()) != {"boxes", "labels"}  # avoid confusion with obj detection target
             ):
                 img_transformed = _copy_tensor(img)
@@ -82,6 +82,7 @@ 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

doctr/datasets/datasets/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -18,6 +18,8 @@ __all__ = ["AbstractDataset", "VisionDataset"]
 
 
 class AbstractDataset(_AbstractDataset):
+    """Abstract class for all datasets"""
+
     def _read_sample(self, index: int) -> Tuple[torch.Tensor, Any]:
         img_name, target = self.data[index]
 
@@ -53,5 +55,5 @@ class AbstractDataset(_AbstractDataset):
         return images, list(targets)
 
 
-class VisionDataset(AbstractDataset, _VisionDataset):
+class VisionDataset(AbstractDataset, _VisionDataset):  # noqa: D101
     pass

doctr/datasets/datasets/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -18,6 +18,8 @@ __all__ = ["AbstractDataset", "VisionDataset"]
 
 
 class AbstractDataset(_AbstractDataset):
+    """Abstract class for all datasets"""
+
     def _read_sample(self, index: int) -> Tuple[tf.Tensor, Any]:
         img_name, target = self.data[index]
 
@@ -53,5 +55,5 @@ class AbstractDataset(_AbstractDataset):
         return images, list(targets)
 
 
-class VisionDataset(AbstractDataset, _VisionDataset):
+class VisionDataset(AbstractDataset, _VisionDataset):  # noqa: D101
     pass

doctr/datasets/detection.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -26,6 +26,7 @@ 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)
@@ -66,14 +67,16 @@ class DetectionDataset(AbstractDataset):
     def format_polygons(
         self, polygons: Union[List, Dict], use_polygons: bool, np_dtype: Type
     ) -> Tuple[np.ndarray, List[str]]:
-        """format polygons into an array
+        """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
         """
@@ -92,4 +95,4 @@ class DetectionDataset(AbstractDataset):
 
     @property
     def class_names(self):
-        return sorted(list(set(self._class_names)))
+        return sorted(set(self._class_names))

doctr/datasets/doc_artefacts.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -26,6 +26,7 @@ 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`.

doctr/datasets/funsd.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -29,6 +29,7 @@ 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
@@ -81,7 +82,7 @@ class FUNSD(VisionDataset):
             text_targets, box_targets = zip(*_targets)
             if use_polygons:
                 # xmin, ymin, xmax, ymax -> (x, y) coordinates of top left, top right, bottom right, bottom left corners
-                box_targets = [
+                box_targets = [  # type: ignore[assignment]
                     [
                         [box[0], box[1]],
                         [box[2], box[1]],
@@ -100,12 +101,10 @@ class FUNSD(VisionDataset):
                     if not any(char in label for char in ["☑", "☐", "\uf703", "\uf702"]):
                         self.data.append((crop, label))
             else:
-                self.data.append(
-                    (
-                        img_path,
-                        dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=list(text_targets)),
-                    )
-                )
+                self.data.append((
+                    img_path,
+                    dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=list(text_targets)),
+                ))
 
         self.root = tmp_root
 

doctr/datasets/generator/base.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -24,6 +24,7 @@ def synthesize_text_img(
     """Generate a synthetic text image
 
     Args:
+    ----
         text: the text to render as an image
         font_size: the size of the font
         font_family: the font family (has to be installed on your system)
@@ -31,9 +32,9 @@ def synthesize_text_img(
         text_color: text color on the final image
 
     Returns:
+    -------
         PIL image of the text
     """
-
     background_color = (0, 0, 0) if background_color is None else background_color
     text_color = (255, 255, 255) if text_color is None else text_color
 

doctr/datasets/generator/pytorch.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -18,6 +18,7 @@ class CharacterGenerator(_CharacterGenerator):
     >>> img, target = ds[0]
 
     Args:
+    ----
         vocab: vocabulary to take the character from
         num_samples: number of samples that will be generated iterating over the dataset
         cache_samples: whether generated images should be cached firsthand
@@ -39,6 +40,7 @@ class WordGenerator(_WordGenerator):
     >>> img, target = ds[0]
 
     Args:
+    ----
         vocab: vocabulary to take the character from
         min_chars: minimum number of characters in a word
         max_chars: maximum number of characters in a word

doctr/datasets/generator/tensorflow.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -18,6 +18,7 @@ class CharacterGenerator(_CharacterGenerator):
     >>> img, target = ds[0]
 
     Args:
+    ----
         vocab: vocabulary to take the character from
         num_samples: number of samples that will be generated iterating over the dataset
         cache_samples: whether generated images should be cached firsthand
@@ -45,6 +46,7 @@ class WordGenerator(_WordGenerator):
     >>> img, target = ds[0]
 
     Args:
+    ----
         vocab: vocabulary to take the character from
         min_chars: minimum number of characters in a word
         max_chars: maximum number of characters in a word

doctr/datasets/ic03.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -28,6 +28,7 @@ class IC03(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
@@ -73,7 +74,7 @@ class IC03(VisionDataset):
         xml_root = xml_tree.getroot()
 
         for image in tqdm(iterable=xml_root, desc="Unpacking IC03", total=len(xml_root)):
-            name, resolution, rectangles = image
+            name, _resolution, rectangles = image
 
             # File existence check
             if not os.path.exists(os.path.join(tmp_root, name.text)):

doctr/datasets/ic13.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -33,6 +33,7 @@ class IC13(AbstractDataset):
     >>> img, target = test_set[0]
 
     Args:
+    ----
         img_folder: folder with all the images of the dataset
         label_folder: folder with all annotation files for the images
         use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)

doctr/datasets/iiit5k.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -30,6 +30,7 @@ class IIIT5K(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
@@ -91,9 +92,10 @@ class IIIT5K(VisionDataset):
                     box_targets = [[box[0], box[1], box[0] + box[2], box[1] + box[3]] for box in box_targets]
 
                 # label are casted to list where each char corresponds to the character's bounding box
-                self.data.append(
-                    (_raw_path, dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=list(_raw_label)))
-                )
+                self.data.append((
+                    _raw_path,
+                    dict(boxes=np.asarray(box_targets, dtype=np_dtype), labels=list(_raw_label)),
+                ))
 
         self.root = tmp_root
 

doctr/datasets/iiithws.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -32,6 +32,7 @@ class IIITHWS(AbstractDataset):
     >>> img, target = test_set[0]
 
     Args:
+    ----
         img_folder: folder with all the images of the dataset
         label_path: path to the file with the labels
         train: whether the subset should be the training one

doctr/datasets/imgur5k.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -40,6 +40,7 @@ class IMGUR5K(AbstractDataset):
     >>> img, target = test_set[0]
 
     Args:
+    ----
         img_folder: folder with all the images of the dataset
         label_path: path to the annotations file of the dataset
         train: whether the subset should be the training one
@@ -111,7 +112,7 @@ class IMGUR5K(AbstractDataset):
                 if ann["word"] != "."
             ]
             # (x, y) coordinates of top left, top right, bottom right, bottom left corners
-            box_targets = [cv2.boxPoints(((box[0], box[1]), (box[2], box[3]), box[4])) for box in _boxes]
+            box_targets = [cv2.boxPoints(((box[0], box[1]), (box[2], box[3]), box[4])) for box in _boxes]  # type: ignore[arg-type]
 
             if not use_polygons:
                 # xmin, ymin, xmax, ymax

doctr/datasets/loader.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -18,12 +18,13 @@ def default_collate(samples):
     """Collate multiple elements into batches
 
     Args:
+    ----
         samples: list of N tuples containing M elements
 
     Returns:
+    -------
         Tuple of M sequences contianing N elements each
     """
-
     batch_data = zip(*samples)
 
     tf_data = tuple(tf.stack(elt, axis=0) for elt in batch_data)
@@ -41,6 +42,7 @@ class DataLoader:
     >>> images, targets = next(train_iter)
 
     Args:
+    ----
         dataset: the dataset
         shuffle: whether the samples should be shuffled before passing it to the iterator
         batch_size: number of elements in each batch

doctr/datasets/mjsynth.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -30,6 +30,7 @@ class MJSynth(AbstractDataset):
     >>> img, target = test_set[0]
 
     Args:
+    ----
         img_folder: folder with all the images of the dataset
         label_path: path to the file with the labels
         train: whether the subset should be the training one

doctr/datasets/ocr.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -24,6 +24,7 @@ class OCRDataset(AbstractDataset):
     >>> img, target = train_set[0]
 
     Args:
+    ----
         img_folder: local path to image folder (all jpg at the root)
         label_file: local path to the label file
         use_polygons: whether polygons should be considered as rotated bounding box (instead of straight ones)

doctr/datasets/orientation.py

@@ -0,0 +1,40 @@
+# Copyright (C) 2021-2024, 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 typing import Any, List, Tuple
+
+import numpy as np
+
+from .datasets import AbstractDataset
+
+__all__ = ["OrientationDataset"]
+
+
+class OrientationDataset(AbstractDataset):
+    """Implements a basic image dataset where targets are filled with zeros.
+
+    >>> from doctr.datasets import OrientationDataset
+    >>> train_set = OrientationDataset(img_folder="/path/to/images")
+    >>> img, target = train_set[0]
+
+    Args:
+    ----
+        img_folder: folder with all the images of the dataset
+        **kwargs: keyword arguments from `AbstractDataset`.
+    """
+
+    def __init__(
+        self,
+        img_folder: str,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            img_folder,
+            **kwargs,
+        )
+
+        # initialize dataset with 0 degree rotation targets
+        self.data: List[Tuple[str, np.ndarray]] = [(img_name, np.array([0])) for img_name in os.listdir(self.root)]

doctr/datasets/recognition.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -22,6 +22,7 @@ class RecognitionDataset(AbstractDataset):
     >>> img, target = train_set[0]
 
     Args:
+    ----
         img_folder: path to the images folder
         labels_path: pathe to the json file containing all labels (character sequences)
         **kwargs: keyword arguments from `AbstractDataset`.
@@ -36,7 +37,7 @@ class RecognitionDataset(AbstractDataset):
         super().__init__(img_folder, **kwargs)
 
         self.data: List[Tuple[str, str]] = []
-        with open(labels_path) as f:
+        with open(labels_path, encoding="utf-8") as f:
             labels = json.load(f)
 
         for img_name, label in labels.items():

doctr/datasets/sroie.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -29,6 +29,7 @@ class SROIE(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

doctr/datasets/svhn.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -28,6 +28,7 @@ class SVHN(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

doctr/datasets/svt.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -28,6 +28,7 @@ class SVT(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
@@ -66,7 +67,7 @@ class SVT(VisionDataset):
         xml_root = xml_tree.getroot()
 
         for image in tqdm(iterable=xml_root, desc="Unpacking SVT", total=len(xml_root)):
-            name, _, _, resolution, rectangles = image
+            name, _, _, _resolution, rectangles = image
 
             # File existence check
             if not os.path.exists(os.path.join(tmp_root, name.text)):

doctr/datasets/synthtext.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021-2023, Mindee.
+# Copyright (C) 2021-2024, 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.
@@ -31,6 +31,7 @@ class SynthText(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