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

doctr/utils/multithreading.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,15 +6,16 @@
 
 import multiprocessing as mp
 import os
+from collections.abc import Callable, Iterable, Iterator
 from multiprocessing.pool import ThreadPool
-from typing import Any, Callable, Iterable, Iterator, Optional
+from typing import Any
 
 from doctr.file_utils import ENV_VARS_TRUE_VALUES
 
 __all__ = ["multithread_exec"]
 
 
-def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Optional[int] = None) -> Iterator[Any]:
+def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: int | None = None) -> Iterator[Any]:
     """Execute a given function in parallel for each element of a given sequence
 
     >>> from doctr.utils.multithreading import multithread_exec
@@ -22,17 +23,14 @@ def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Op
     >>> results = multithread_exec(lambda x: x ** 2, entries)
 
     Args:
-    ----
         func: function to be executed on each element of the iterable
         seq: iterable
         threads: number of workers to be used for multiprocessing
 
     Returns:
-    -------
         iterator of the function's results using the iterable as inputs
 
     Notes:
-    -----
         This function uses ThreadPool from multiprocessing package, which uses `/dev/shm` directory for shared memory.
         If you do not have write permissions for this directory (if you run `doctr` on AWS Lambda for instance),
         you might want to disable multiprocessing. To achieve that, set 'DOCTR_MULTIPROCESSING_DISABLE' to 'TRUE'.

doctr/utils/reconstitution.py

@@ -1,8 +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, Optional
+import logging
+from typing import Any
 
 import numpy as np
 from anyascii import anyascii
@@ -13,114 +14,195 @@ from .fonts import get_font
 __all__ = ["synthesize_page", "synthesize_kie_page"]
 
 
+# Global variable to avoid multiple warnings
+ROTATION_WARNING = False
+
+
+def _warn_rotation(entry: dict[str, Any]) -> None:  # pragma: no cover
+    global ROTATION_WARNING
+    if not ROTATION_WARNING and len(entry["geometry"]) == 4:
+        logging.warning("Polygons with larger rotations will lead to inaccurate rendering")
+        ROTATION_WARNING = True
+
+
+def _synthesize(
+    response: Image.Image,
+    entry: dict[str, Any],
+    w: int,
+    h: int,
+    draw_proba: bool = False,
+    font_family: str | None = None,
+    smoothing_factor: float = 0.75,
+    min_font_size: int = 6,
+    max_font_size: int = 50,
+) -> Image.Image:
+    if len(entry["geometry"]) == 2:
+        (xmin, ymin), (xmax, ymax) = entry["geometry"]
+        polygon = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
+    else:
+        polygon = entry["geometry"]
+
+    # Calculate the bounding box of the word
+    x_coords, y_coords = zip(*polygon)
+    xmin, ymin, xmax, ymax = (
+        int(round(w * min(x_coords))),
+        int(round(h * min(y_coords))),
+        int(round(w * max(x_coords))),
+        int(round(h * max(y_coords))),
+    )
+    word_width = xmax - xmin
+    word_height = ymax - ymin
+
+    # If lines are provided instead of words, concatenate the word entries
+    if "words" in entry:
+        word_text = " ".join(word["value"] for word in entry["words"])
+    else:
+        word_text = entry["value"]
+    # Find the optimal font size
+    try:
+        font_size = min(word_height, max_font_size)
+        font = get_font(font_family, font_size)
+        text_width, text_height = font.getbbox(word_text)[2:4]
+
+        while (text_width > word_width or text_height > word_height) and font_size > min_font_size:
+            font_size = max(int(font_size * smoothing_factor), min_font_size)
+            font = get_font(font_family, font_size)
+            text_width, text_height = font.getbbox(word_text)[2:4]
+    except ValueError:
+        font = get_font(font_family, min_font_size)
+
+    # Create a mask for the word
+    mask = Image.new("L", (w, h), 0)
+    ImageDraw.Draw(mask).polygon([(int(round(w * x)), int(round(h * y))) for x, y in polygon], fill=255)
+
+    # Draw the word text
+    d = ImageDraw.Draw(response)
+    try:
+        try:
+            d.text((xmin, ymin), word_text, font=font, fill=(0, 0, 0), anchor="lt")
+        except UnicodeEncodeError:
+            d.text((xmin, ymin), anyascii(word_text), font=font, fill=(0, 0, 0), anchor="lt")
+    # Catch generic exceptions to avoid crashing the whole rendering
+    except Exception:  # pragma: no cover
+        logging.warning(f"Could not render word: {word_text}")
+
+    if draw_proba:
+        confidence = (
+            entry["confidence"]
+            if "confidence" in entry
+            else sum(w["confidence"] for w in entry["words"]) / len(entry["words"])
+        )
+        p = int(255 * confidence)
+        color = (255 - p, 0, p)  # Red to blue gradient based on probability
+        d.rectangle([(xmin, ymin), (xmax, ymax)], outline=color, width=2)
+
+        prob_font = get_font(font_family, 20)
+        prob_text = f"{confidence:.2f}"
+        prob_text_width, prob_text_height = prob_font.getbbox(prob_text)[2:4]
+
+        # Position the probability slightly above the bounding box
+        prob_x_offset = (word_width - prob_text_width) // 2
+        prob_y_offset = ymin - prob_text_height - 2
+        prob_y_offset = max(0, prob_y_offset)
+
+        d.text((xmin + prob_x_offset, prob_y_offset), prob_text, font=prob_font, fill=color, anchor="lt")
+
+    return response
+
+
 def synthesize_page(
-    page: Dict[str, Any],
+    page: dict[str, Any],
     draw_proba: bool = False,
-    font_family: Optional[str] = None,
+    font_family: str | None = None,
+    smoothing_factor: float = 0.95,
+    min_font_size: int = 8,
+    max_font_size: int = 50,
 ) -> np.ndarray:
     """Draw a the content of the element page (OCR response) on a blank page.
 
     Args:
-    ----
         page: exported Page object to represent
         draw_proba: if True, draw words in colors to represent confidence. Blue: p=1, red: p=0
-        font_size: size of the font, default font = 13
         font_family: family of the font
+        smoothing_factor: factor to smooth the font size
+        min_font_size: minimum font size
+        max_font_size: maximum font size
 
     Returns:
-    -------
         the synthesized page
     """
     # Draw template
     h, w = page["dimensions"]
-    response = 255 * np.ones((h, w, 3), dtype=np.int32)
+    response = Image.new("RGB", (w, h), color=(255, 255, 255))
 
-    # Draw each word
     for block in page["blocks"]:
-        for line in block["lines"]:
-            for word in line["words"]:
-                # Get absolute word geometry
-                (xmin, ymin), (xmax, ymax) = word["geometry"]
-                xmin, xmax = int(round(w * xmin)), int(round(w * xmax))
-                ymin, ymax = int(round(h * ymin)), int(round(h * ymax))
-
-                # White drawing context adapted to font size, 0.75 factor to convert pts --> pix
-                font = get_font(font_family, int(0.75 * (ymax - ymin)))
-                img = Image.new("RGB", (xmax - xmin, ymax - ymin), color=(255, 255, 255))
-                d = ImageDraw.Draw(img)
-                # Draw in black the value of the word
-                try:
-                    d.text((0, 0), word["value"], font=font, fill=(0, 0, 0))
-                except UnicodeEncodeError:
-                    # When character cannot be encoded, use its anyascii version
-                    d.text((0, 0), anyascii(word["value"]), font=font, fill=(0, 0, 0))
-
-                # Colorize if draw_proba
-                if draw_proba:
-                    p = int(255 * word["confidence"])
-                    mask = np.where(np.array(img) == 0, 1, 0)
-                    proba: np.ndarray = np.array([255 - p, 0, p])
-                    color = mask * proba[np.newaxis, np.newaxis, :]
-                    white_mask = 255 * (1 - mask)
-                    img = color + white_mask
-
-                # Write to response page
-                response[ymin:ymax, xmin:xmax, :] = np.array(img)
-
-    return response
+        # If lines are provided use these to get better rendering results
+        if len(block["lines"]) > 1:
+            for line in block["lines"]:
+                _warn_rotation(block)  # pragma: no cover
+                response = _synthesize(
+                    response=response,
+                    entry=line,
+                    w=w,
+                    h=h,
+                    draw_proba=draw_proba,
+                    font_family=font_family,
+                    smoothing_factor=smoothing_factor,
+                    min_font_size=min_font_size,
+                    max_font_size=max_font_size,
+                )
+        # Otherwise, draw each word
+        else:
+            for line in block["lines"]:
+                _warn_rotation(block)  # pragma: no cover
+                for word in line["words"]:
+                    response = _synthesize(
+                        response=response,
+                        entry=word,
+                        w=w,
+                        h=h,
+                        draw_proba=draw_proba,
+                        font_family=font_family,
+                        smoothing_factor=smoothing_factor,
+                        min_font_size=min_font_size,
+                        max_font_size=max_font_size,
+                    )
+
+    return np.array(response, dtype=np.uint8)
 
 
 def synthesize_kie_page(
-    page: Dict[str, Any],
+    page: dict[str, Any],
     draw_proba: bool = False,
-    font_family: Optional[str] = None,
+    font_family: str | None = None,
 ) -> np.ndarray:
     """Draw a the content of the element page (OCR response) on a blank page.
 
     Args:
-    ----
         page: exported Page object to represent
         draw_proba: if True, draw words in colors to represent confidence. Blue: p=1, red: p=0
-        font_size: size of the font, default font = 13
         font_family: family of the font
+        smoothing_factor: factor to smooth the font size
+        min_font_size: minimum font size
+        max_font_size: maximum font size
 
     Returns:
-    -------
         the synthesized page
     """
     # Draw template
     h, w = page["dimensions"]
-    response = 255 * np.ones((h, w, 3), dtype=np.int32)
+    response = Image.new("RGB", (w, h), color=(255, 255, 255))
 
     # Draw each word
     for predictions in page["predictions"].values():
         for prediction in predictions:
-            # Get aboslute word geometry
-            (xmin, ymin), (xmax, ymax) = prediction["geometry"]
-            xmin, xmax = int(round(w * xmin)), int(round(w * xmax))
-            ymin, ymax = int(round(h * ymin)), int(round(h * ymax))
-
-            # White drawing context adapted to font size, 0.75 factor to convert pts --> pix
-            font = get_font(font_family, int(0.75 * (ymax - ymin)))
-            img = Image.new("RGB", (xmax - xmin, ymax - ymin), color=(255, 255, 255))
-            d = ImageDraw.Draw(img)
-            # Draw in black the value of the word
-            try:
-                d.text((0, 0), prediction["value"], font=font, fill=(0, 0, 0))
-            except UnicodeEncodeError:
-                # When character cannot be encoded, use its anyascii version
-                d.text((0, 0), anyascii(prediction["value"]), font=font, fill=(0, 0, 0))
-
-            # Colorize if draw_proba
-            if draw_proba:
-                p = int(255 * prediction["confidence"])
-                mask = np.where(np.array(img) == 0, 1, 0)
-                proba: np.ndarray = np.array([255 - p, 0, p])
-                color = mask * proba[np.newaxis, np.newaxis, :]
-                white_mask = 255 * (1 - mask)
-                img = color + white_mask
-
-            # Write to response page
-            response[ymin:ymax, xmin:xmax, :] = np.array(img)
-
-    return response
+            _warn_rotation(prediction)  # pragma: no cover
+            response = _synthesize(
+                response=response,
+                entry=prediction,
+                w=w,
+                h=h,
+                draw_proba=draw_proba,
+                font_family=font_family,
+            )
+    return np.array(response, dtype=np.uint8)

doctr/utils/repr.py

@@ -1,11 +1,10 @@
-# 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.
 
 # Adapted from https://github.com/pytorch/torch/blob/master/torch/nn/modules/module.py
 
-from typing import List
 
 __all__ = ["NestedObject"]
 
@@ -25,7 +24,7 @@ def _addindent(s_, num_spaces):
 class NestedObject:
     """Base class for all nested objects in doctr"""
 
-    _children_names: List[str]
+    _children_names: list[str]
 
     def extra_repr(self) -> str:
         return ""

doctr/utils/visualization.py

@@ -1,10 +1,10 @@
-# 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 colorsys
 from copy import deepcopy
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 import cv2
 import matplotlib.patches as patches
@@ -19,9 +19,9 @@ __all__ = ["visualize_page", "visualize_kie_page", "draw_boxes"]
 
 def rect_patch(
     geometry: BoundingBox,
-    page_dimensions: Tuple[int, int],
-    label: Optional[str] = None,
-    color: Tuple[float, float, float] = (0, 0, 0),
+    page_dimensions: tuple[int, int],
+    label: str | None = None,
+    color: tuple[float, float, float] = (0, 0, 0),
     alpha: float = 0.3,
     linewidth: int = 2,
     fill: bool = True,
@@ -30,7 +30,6 @@ def rect_patch(
     """Create a matplotlib rectangular patch for the element
 
     Args:
-    ----
         geometry: bounding box of the element
         page_dimensions: dimensions of the Page in format (height, width)
         label: label to display when hovered
@@ -41,7 +40,6 @@ def rect_patch(
         preserve_aspect_ratio: pass True if you passed True to the predictor
 
     Returns:
-    -------
         a rectangular Patch
     """
     if len(geometry) != 2 or any(not isinstance(elt, tuple) or len(elt) != 2 for elt in geometry):
@@ -70,9 +68,9 @@ def rect_patch(
 
 def polygon_patch(
     geometry: np.ndarray,
-    page_dimensions: Tuple[int, int],
-    label: Optional[str] = None,
-    color: Tuple[float, float, float] = (0, 0, 0),
+    page_dimensions: tuple[int, int],
+    label: str | None = None,
+    color: tuple[float, float, float] = (0, 0, 0),
     alpha: float = 0.3,
     linewidth: int = 2,
     fill: bool = True,
@@ -81,7 +79,6 @@ def polygon_patch(
     """Create a matplotlib polygon patch for the element
 
     Args:
-    ----
         geometry: bounding box of the element
         page_dimensions: dimensions of the Page in format (height, width)
         label: label to display when hovered
@@ -92,7 +89,6 @@ def polygon_patch(
         preserve_aspect_ratio: pass True if you passed True to the predictor
 
     Returns:
-    -------
         a polygon Patch
     """
     if not geometry.shape == (4, 2):
@@ -114,20 +110,18 @@ def polygon_patch(
 
 
 def create_obj_patch(
-    geometry: Union[BoundingBox, Polygon4P, np.ndarray],
-    page_dimensions: Tuple[int, int],
+    geometry: BoundingBox | Polygon4P | np.ndarray,
+    page_dimensions: tuple[int, int],
     **kwargs: Any,
 ) -> patches.Patch:
     """Create a matplotlib patch for the element
 
     Args:
-    ----
         geometry: bounding box (straight or rotated) of the element
         page_dimensions: dimensions of the page in format (height, width)
         **kwargs: keyword arguments for the patch
 
     Returns:
-    -------
         a matplotlib Patch
     """
     if isinstance(geometry, tuple):
@@ -140,15 +134,13 @@ def create_obj_patch(
     raise ValueError("invalid geometry format")
 
 
-def get_colors(num_colors: int) -> List[Tuple[float, float, float]]:
+def get_colors(num_colors: int) -> list[tuple[float, float, float]]:
     """Generate num_colors color for matplotlib
 
     Args:
-    ----
         num_colors: number of colors to generate
 
     Returns:
-    -------
         colors: list of generated colors
     """
     colors = []
@@ -161,7 +153,7 @@ def get_colors(num_colors: int) -> List[Tuple[float, float, float]]:
 
 
 def visualize_page(
-    page: Dict[str, Any],
+    page: dict[str, Any],
     image: np.ndarray,
     words_only: bool = True,
     display_artefacts: bool = True,
@@ -183,7 +175,6 @@ def visualize_page(
     >>> plt.show()
 
     Args:
-    ----
         page: the exported Page of a Document
         image: np array of the page, needs to have the same shape than page['dimensions']
         words_only: whether only words should be displayed
@@ -194,7 +185,6 @@ def visualize_page(
         **kwargs: keyword arguments for the polygon patch
 
     Returns:
-    -------
         the matplotlib figure
     """
     # Get proper scale and aspect ratio
@@ -207,7 +197,7 @@ def visualize_page(
     ax.axis("off")
 
     if interactive:
-        artists: List[patches.Patch] = []  # instantiate an empty list of patches (to be drawn on the page)
+        artists: list[patches.Patch] = []  # instantiate an empty list of patches (to be drawn on the page)
 
     for block in page["blocks"]:
         if not words_only:
@@ -287,7 +277,7 @@ def visualize_page(
 
 
 def visualize_kie_page(
-    page: Dict[str, Any],
+    page: dict[str, Any],
     image: np.ndarray,
     words_only: bool = False,
     display_artefacts: bool = True,
@@ -309,7 +299,6 @@ def visualize_kie_page(
     >>> plt.show()
 
     Args:
-    ----
         page: the exported Page of a Document
         image: np array of the page, needs to have the same shape than page['dimensions']
         words_only: whether only words should be displayed
@@ -320,7 +309,6 @@ def visualize_kie_page(
         **kwargs: keyword arguments for the polygon patch
 
     Returns:
-    -------
         the matplotlib figure
     """
     # Get proper scale and aspect ratio
@@ -333,7 +321,7 @@ def visualize_kie_page(
     ax.axis("off")
 
     if interactive:
-        artists: List[patches.Patch] = []  # instantiate an empty list of patches (to be drawn on the page)
+        artists: list[patches.Patch] = []  # instantiate an empty list of patches (to be drawn on the page)
 
     colors = {k: color for color, k in zip(get_colors(len(page["predictions"])), page["predictions"])}
     for key, value in page["predictions"].items():
@@ -363,11 +351,10 @@ def visualize_kie_page(
     return fig
 
 
-def draw_boxes(boxes: np.ndarray, image: np.ndarray, color: Optional[Tuple[int, int, int]] = None, **kwargs) -> None:
+def draw_boxes(boxes: np.ndarray, image: np.ndarray, color: tuple[int, int, int] | None = None, **kwargs) -> None:
     """Draw an array of relative straight boxes on an image
 
     Args:
-    ----
         boxes: array of relative boxes, of shape (*, 4)
         image: np array, float32 or uint8
         color: color to use for bounding box edges

doctr/version.py

@@ -1 +1 @@
-__version__ = 'v0.9.0'
+__version__ = 'v0.11.0'