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

doctr/io/image/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.
@@ -20,13 +20,14 @@ def tensor_from_pil(pil_img: Image, dtype: torch.dtype = torch.float32) -> torch
     """Convert a PIL Image to a PyTorch tensor
 
     Args:
+    ----
         pil_img: a PIL image
         dtype: the output tensor data type
 
     Returns:
+    -------
         decoded image as tensor
     """
-
     if dtype == torch.float32:
         img = to_tensor(pil_img)
     else:
@@ -39,13 +40,14 @@ def read_img_as_tensor(img_path: AbstractPath, dtype: torch.dtype = torch.float3
     """Read an image file as a PyTorch tensor
 
     Args:
+    ----
         img_path: location of the image file
         dtype: the desired data type of the output tensor. If it is float-related, values will be divided by 255.
 
     Returns:
+    -------
         decoded image as a tensor
     """
-
     if dtype not in (torch.uint8, torch.float16, torch.float32):
         raise ValueError("insupported value for dtype")
 
@@ -58,13 +60,14 @@ def decode_img_as_tensor(img_content: bytes, dtype: torch.dtype = torch.float32)
     """Read a byte stream as a PyTorch tensor
 
     Args:
+    ----
         img_content: bytes of a decoded image
         dtype: the desired data type of the output tensor. If it is float-related, values will be divided by 255.
 
     Returns:
+    -------
         decoded image as a tensor
     """
-
     if dtype not in (torch.uint8, torch.float16, torch.float32):
         raise ValueError("insupported value for dtype")
 
@@ -77,13 +80,14 @@ def tensor_from_numpy(npy_img: np.ndarray, dtype: torch.dtype = torch.float32) -
     """Read an image file as a PyTorch tensor
 
     Args:
-        img: image encoded as a numpy array of shape (H, W, C) in np.uint8
+    ----
+        npy_img: image encoded as a numpy array of shape (H, W, C) in np.uint8
         dtype: the desired data type of the output tensor. If it is float-related, values will be divided by 255.
 
     Returns:
+    -------
         same image as a tensor of shape (C, H, W)
     """
-
     if dtype not in (torch.uint8, torch.float16, torch.float32):
         raise ValueError("insupported value for dtype")
 
@@ -101,4 +105,5 @@ def tensor_from_numpy(npy_img: np.ndarray, dtype: torch.dtype = torch.float32) -
 
 
 def get_img_shape(img: torch.Tensor) -> Tuple[int, int]:
-    return img.shape[-2:]  # type: ignore[return-value]
+    """Get the shape of an image"""
+    return img.shape[-2:]

doctr/io/image/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.
@@ -19,13 +19,14 @@ def tensor_from_pil(pil_img: Image, dtype: tf.dtypes.DType = tf.float32) -> tf.T
     """Convert a PIL Image to a TensorFlow tensor
 
     Args:
+    ----
         pil_img: a PIL image
         dtype: the output tensor data type
 
     Returns:
+    -------
         decoded image as tensor
     """
-
     npy_img = img_to_array(pil_img)
 
     return tensor_from_numpy(npy_img, dtype)
@@ -35,13 +36,14 @@ def read_img_as_tensor(img_path: AbstractPath, dtype: tf.dtypes.DType = tf.float
     """Read an image file as a TensorFlow tensor
 
     Args:
+    ----
         img_path: location of the image file
         dtype: the desired data type of the output tensor. If it is float-related, values will be divided by 255.
 
     Returns:
+    -------
         decoded image as a tensor
     """
-
     if dtype not in (tf.uint8, tf.float16, tf.float32):
         raise ValueError("insupported value for dtype")
 
@@ -59,13 +61,14 @@ def decode_img_as_tensor(img_content: bytes, dtype: tf.dtypes.DType = tf.float32
     """Read a byte stream as a TensorFlow tensor
 
     Args:
+    ----
         img_content: bytes of a decoded image
         dtype: the desired data type of the output tensor. If it is float-related, values will be divided by 255.
 
     Returns:
+    -------
         decoded image as a tensor
     """
-
     if dtype not in (tf.uint8, tf.float16, tf.float32):
         raise ValueError("insupported value for dtype")
 
@@ -82,13 +85,14 @@ def tensor_from_numpy(npy_img: np.ndarray, dtype: tf.dtypes.DType = tf.float32)
     """Read an image file as a TensorFlow tensor
 
     Args:
-        img: image encoded as a numpy array of shape (H, W, C) in np.uint8
+    ----
+        npy_img: image encoded as a numpy array of shape (H, W, C) in np.uint8
         dtype: the desired data type of the output tensor. If it is float-related, values will be divided by 255.
 
     Returns:
+    -------
         same image as a tensor of shape (H, W, C)
     """
-
     if dtype not in (tf.uint8, tf.float16, tf.float32):
         raise ValueError("insupported value for dtype")
 
@@ -102,4 +106,5 @@ def tensor_from_numpy(npy_img: np.ndarray, dtype: tf.dtypes.DType = tf.float32)
 
 
 def get_img_shape(img: tf.Tensor) -> Tuple[int, int]:
+    """Get the shape of an image"""
     return img.shape[:2]

doctr/io/pdf.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,20 +22,21 @@ def read_pdf(
 ) -> List[np.ndarray]:
     """Read a PDF file and convert it into an image in numpy format
 
-    >>> from doctr.documents import read_pdf
+    >>> from doctr.io import read_pdf
     >>> doc = read_pdf("path/to/your/doc.pdf")
 
     Args:
+    ----
         file: the path to the PDF file
         scale: rendering scale (1 corresponds to 72dpi)
         rgb_mode: if True, the output will be RGB, otherwise BGR
         password: a password to unlock the document, if encrypted
-        kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`
+        **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`
 
     Returns:
+    -------
         the list of pages decoded as numpy ndarray of shape H x W x C
     """
-
     # Rasterise pages to numpy ndarrays with pypdfium2
     pdf = pdfium.PdfDocument(file, password=password, autoclose=True)
     return [page.render(scale=scale, rev_byteorder=rgb_mode, **kwargs).to_numpy() for page in pdf]

doctr/io/reader.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,29 +24,34 @@ class DocumentFile:
     def from_pdf(cls, file: AbstractFile, **kwargs) -> List[np.ndarray]:
         """Read a PDF file
 
-        >>> from doctr.documents import DocumentFile
+        >>> from doctr.io import DocumentFile
         >>> doc = DocumentFile.from_pdf("path/to/your/doc.pdf")
 
         Args:
+        ----
             file: the path to the PDF file or a binary stream
+            **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`
 
         Returns:
+        -------
             the list of pages decoded as numpy ndarray of shape H x W x 3
         """
-
         return read_pdf(file, **kwargs)
 
     @classmethod
     def from_url(cls, url: str, **kwargs) -> List[np.ndarray]:
         """Interpret a web page as a PDF document
 
-        >>> from doctr.documents import DocumentFile
+        >>> from doctr.io import DocumentFile
         >>> doc = DocumentFile.from_url("https://www.yoursite.com")
 
         Args:
+        ----
             url: the URL of the target web page
+            **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render`
 
         Returns:
+        -------
             the list of pages decoded as numpy ndarray of shape H x W x 3
         """
         pdf_stream = read_html(url)
@@ -56,13 +61,16 @@ class DocumentFile:
     def from_images(cls, files: Union[Sequence[AbstractFile], AbstractFile], **kwargs) -> List[np.ndarray]:
         """Read an image file (or a collection of image files) and convert it into an image in numpy format
 
-        >>> from doctr.documents import DocumentFile
+        >>> from doctr.io import DocumentFile
         >>> pages = DocumentFile.from_images(["path/to/your/page1.png", "path/to/your/page2.png"])
 
         Args:
+        ----
             files: the path to the image file or a binary stream, or a collection of those
+            **kwargs: additional parameters to :meth:`doctr.io.image.read_img_as_numpy`
 
         Returns:
+        -------
             the list of pages decoded as numpy ndarray of shape H x W x 3
         """
         if isinstance(files, (str, Path, bytes)):

doctr/models/_utils.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.
@@ -11,43 +11,54 @@ import cv2
 import numpy as np
 from langdetect import LangDetectException, detect_langs
 
-__all__ = ["estimate_orientation", "get_bitmap_angle", "get_language", "invert_data_structure"]
+__all__ = ["estimate_orientation", "get_language", "invert_data_structure"]
 
 
 def get_max_width_length_ratio(contour: np.ndarray) -> float:
     """Get the maximum shape ratio of a contour.
 
     Args:
+    ----
         contour: the contour from cv2.findContour
 
-    Returns: the maximum shape ratio
+    Returns:
+    -------
+        the maximum shape ratio
     """
     _, (w, h), _ = cv2.minAreaRect(contour)
     return max(w / h, h / w)
 
 
-def estimate_orientation(img: np.ndarray, n_ct: int = 50, ratio_threshold_for_lines: float = 5) -> float:
+def estimate_orientation(img: np.ndarray, n_ct: int = 50, ratio_threshold_for_lines: float = 5) -> int:
     """Estimate the angle of the general document orientation based on the
      lines of the document and the assumption that they should be horizontal.
 
     Args:
-        img: the img to analyze
+    ----
+        img: the img or bitmap to analyze (H, W, C)
         n_ct: the number of contours used for the orientation estimation
         ratio_threshold_for_lines: this is the ratio w/h used to discriminates lines
 
     Returns:
+    -------
         the angle of the general document orientation
     """
-    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    gray_img = cv2.medianBlur(gray_img, 5)
-    thresh = cv2.threshold(gray_img, thresh=0, maxval=255, type=cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
+    assert len(img.shape) == 3 and img.shape[-1] in [1, 3], f"Image shape {img.shape} not supported"
+    max_value = np.max(img)
+    min_value = np.min(img)
+    if max_value <= 1 and min_value >= 0 or (max_value <= 255 and min_value >= 0 and img.shape[-1] == 1):
+        thresh = img.astype(np.uint8)
+    if max_value <= 255 and min_value >= 0 and img.shape[-1] == 3:
+        gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        gray_img = cv2.medianBlur(gray_img, 5)
+        thresh = cv2.threshold(gray_img, thresh=0, maxval=255, type=cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]  # type: ignore[assignment]
 
     # try to merge words in lines
     (h, w) = img.shape[:2]
     k_x = max(1, (floor(w / 100)))
     k_y = max(1, (floor(h / 100)))
     kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k_x, k_y))
-    thresh = cv2.dilate(thresh, kernel, iterations=1)
+    thresh = cv2.dilate(thresh, kernel, iterations=1)  # type: ignore[assignment]
 
     # extract contours
     contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
@@ -66,45 +77,8 @@ def estimate_orientation(img: np.ndarray, n_ct: int = 50, ratio_threshold_for_li
     if len(angles) == 0:
         return 0  # in case no angles is found
     else:
-        return -median_low(angles)
-
-
-def get_bitmap_angle(bitmap: np.ndarray, n_ct: int = 20, std_max: float = 3.0) -> float:
-    """From a binarized segmentation map, find contours and fit min area rectangles to determine page angle
-
-    Args:
-        bitmap: binarized segmentation map
-        n_ct: number of contours to use to fit page angle
-        std_max: maximum deviation of the angle distribution to consider the mean angle reliable
-
-    Returns:
-        The angle of the page
-    """
-    # Find all contours on binarized seg map
-    contours, _ = cv2.findContours(bitmap.astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
-    # Sort contours
-    contours = sorted(contours, key=cv2.contourArea, reverse=True)
-
-    # Find largest contours and fit angles
-    # Track heights and widths to find aspect ratio (determine is rotation is clockwise)
-    angles, heights, widths = [], [], []
-    for ct in contours[:n_ct]:
-        _, (w, h), alpha = cv2.minAreaRect(ct)
-        widths.append(w)
-        heights.append(h)
-        angles.append(alpha)
-
-    if np.std(angles) > std_max:
-        # Edge case with angles of both 0 and 90°, or multi_oriented docs
-        angle = 0.0
-    else:
-        angle = -np.mean(angles)
-        # Determine rotation direction (clockwise/counterclockwise)
-        # Angle coverage: [-90°, +90°], half of the quadrant
-        if np.sum(widths) < np.sum(heights):  # CounterClockwise
-            angle = 90 + angle
-
-    return angle
+        median = -median_low(angles)
+        return round(median) if abs(median) != 0 else 0
 
 
 def rectify_crops(
@@ -149,9 +123,13 @@ def rectify_loc_preds(
 def get_language(text: str) -> Tuple[str, float]:
     """Get languages of a text using langdetect model.
     Get the language with the highest probability or no language if only a few words or a low probability
+
     Args:
+    ----
         text (str): text
+
     Returns:
+    -------
         The detected language in ISO 639 code and confidence score
     """
     try:
@@ -164,21 +142,20 @@ def get_language(text: str) -> Tuple[str, float]:
 
 
 def invert_data_structure(
-    x: Union[List[Dict[str, Any]], Dict[str, List[Any]]]
+    x: Union[List[Dict[str, Any]], Dict[str, List[Any]]],
 ) -> Union[List[Dict[str, Any]], Dict[str, List[Any]]]:
     """Invert a List of Dict of elements to a Dict of list of elements and the other way around
 
     Args:
+    ----
         x: a list of dictionaries with the same keys or a dictionary of lists of the same length
 
     Returns:
+    -------
         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(set([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 dictionnary 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/artefacts/barcode.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.
@@ -12,11 +12,11 @@ __all__ = ["BarCodeDetector"]
 
 
 class BarCodeDetector:
-
     """Implements a Bar-code detector.
     For now, only horizontal (or with a small angle) bar-codes are supported
 
     Args:
+    ----
         min_size: minimum relative size of a barcode on the page
         canny_minval: lower bound for canny hysteresis
         canny_maxval: upper-bound for canny hysteresis
@@ -35,7 +35,8 @@ class BarCodeDetector:
         Args:
             img: np image
 
-        Returns:
+        Returns
+        -------
             A list of tuples: [(xmin, ymin, xmax, ymax), ...] containing barcodes rel. coordinates
         """
         # get image size and define parameters

doctr/models/artefacts/face.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.
@@ -14,11 +14,11 @@ __all__ = ["FaceDetector"]
 
 
 class FaceDetector(NestedObject):
-
     """Implements a face detector to detect profile pictures on resumes, IDS, driving licenses, passports...
     Based on open CV CascadeClassifier (haarcascades)
 
     Args:
+    ----
         n_faces: maximal number of faces to detect on a single image, default = 1
     """
 
@@ -42,9 +42,11 @@ class FaceDetector(NestedObject):
         """Detect n_faces on the img
 
         Args:
+        ----
             img: image to detect faces on
 
         Returns:
+        -------
             A list of size n_faces, each face is a tuple of relative xmin, ymin, xmax, ymax
         """
         height, width = img.shape[:2]