[misc] Add minor improvements (#2105)

Author: felixdittrich92

doctr/io/elements.py

@@ -39,6 +39,37 @@
 ]
 
 
+def _resolve_hocr_language(language: dict[str, Any]) -> str:
+    """Resolve the language code to use in the hOCR export, falling back to 'en'.
+
+    Args:
+        language: the page language dictionary `{"value": str | None, "confidence": float | None}`
+
+    Returns:
+        the detected language code when available, 'en' otherwise
+    """
+    lang_value = language.get("value") if isinstance(language, dict) else None
+    return lang_value if isinstance(lang_value, str) and len(lang_value) > 0 else "en"
+
+
+def _hocr_bbox(geometry: BoundingBox, width: int, height: int) -> str:
+    """Format a relative straight bounding box as an absolute hOCR `bbox` property string.
+
+    Args:
+        geometry: the relative bounding box ((xmin, ymin), (xmax, ymax))
+        width: the page width in pixels
+        height: the page height in pixels
+
+    Returns:
+        the hOCR `bbox` property string
+    """
+    (xmin, ymin), (xmax, ymax) = geometry
+    return (
+        f"bbox {int(round(xmin * width))} {int(round(ymin * height))} "
+        f"{int(round(xmax * width))} {int(round(ymax * height))}"
+    )
+
+
 class Element(NestedObject):
     """Implements an abstract document element with exporting and text rendering capabilities"""
 
@@ -518,7 +549,7 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
         line_count: int = 1
         word_count: int = 1
         height, width = self.dimensions
-        language = self.language if "language" in self.language.keys() else "en"
+        language = _resolve_hocr_language(self.language)
         # Create the XML root element
         page_hocr = ETElement("html", attrib={"xmlns": "http://www.w3.org/1999/xhtml", "xml:lang": str(language)})
         # Create the header / SubElements of the root element
@@ -550,15 +581,14 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
         for block in self.blocks:
             if len(block.geometry) != 2:
                 raise TypeError("XML export is only available for straight bounding boxes for now.")
-            (xmin, ymin), (xmax, ymax) = block.geometry
+            block_bbox = _hocr_bbox(block.geometry, width, height)  # type: ignore[arg-type]
             block_div = SubElement(
                 page_div,
                 "div",
                 attrib={
                     "class": "ocr_carea",
                     "id": f"block_{block_count}",
-                    "title": f"bbox {int(round(xmin * width))} {int(round(ymin * height))} \
-                    {int(round(xmax * width))} {int(round(ymax * height))}",
+                    "title": block_bbox,
                 },
             )
             paragraph = SubElement(
@@ -567,38 +597,37 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
                 attrib={
                     "class": "ocr_par",
                     "id": f"par_{block_count}",
-                    "title": f"bbox {int(round(xmin * width))} {int(round(ymin * height))} \
-                    {int(round(xmax * width))} {int(round(ymax * height))}",
+                    "title": block_bbox,
                 },
             )
             block_count += 1
             for line in block.lines:
-                (xmin, ymin), (xmax, ymax) = line.geometry
                 # NOTE: baseline, x_size, x_descenders, x_ascenders is currently initalized to 0
                 line_span = SubElement(
                     paragraph,
                     "span",
                     attrib={
                         "class": "ocr_line",
                         "id": f"line_{line_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",
+                        "title": (
+                            f"{_hocr_bbox(line.geometry, width, height)}; "  # type: ignore[arg-type]
+                            "baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0"
+                        ),
                     },
                 )
                 line_count += 1
                 for word in line.words:
-                    (xmin, ymin), (xmax, ymax) = word.geometry
                     conf = word.confidence
                     word_div = SubElement(
                         line_span,
                         "span",
                         attrib={
                             "class": "ocrx_word",
                             "id": f"word_{word_count}",
-                            "title": f"bbox {int(round(xmin * width))} {int(round(ymin * height))} \
-                            {int(round(xmax * width))} {int(round(ymax * height))}; \
-                            x_wconf {int(round(conf * 100))}",
+                            "title": (
+                                f"{_hocr_bbox(word.geometry, width, height)}; "  # type: ignore[arg-type]
+                                f"x_wconf {int(round(conf * 100))}"
+                            ),
                         },
                     )
                     # set the text
@@ -708,7 +737,7 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
         p_idx = self.page_idx
         prediction_count: int = 1
         height, width = self.dimensions
-        language = self.language if "language" in self.language.keys() else "en"
+        language = _resolve_hocr_language(self.language)
         # Create the XML root element
         page_hocr = ETElement("html", attrib={"xmlns": "http://www.w3.org/1999/xhtml", "xml:lang": str(language)})
         # Create the header / SubElements of the root element
@@ -741,15 +770,14 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
             for prediction in predictions:
                 if len(prediction.geometry) != 2:
                     raise TypeError("XML export is only available for straight bounding boxes for now.")
-                (xmin, ymin), (xmax, ymax) = prediction.geometry
+                prediction_bbox = _hocr_bbox(prediction.geometry, width, height)  # type: ignore[arg-type]
                 prediction_div = SubElement(
                     body,
                     "div",
                     attrib={
                         "class": "ocr_carea",
                         "id": f"{class_name}_prediction_{prediction_count}",
-                        "title": f"bbox {int(round(xmin * width))} {int(round(ymin * height))} \
-                        {int(round(xmax * width))} {int(round(ymax * height))}",
+                        "title": prediction_bbox,
                     },
                 )
                 # NOTE: ocr_par, ocr_line and ocrx_word are the same because the KIE predictions contain only words
@@ -760,8 +788,7 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
                     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))}",
+                        "title": prediction_bbox,
                     },
                 )
                 line_span = SubElement(
@@ -770,9 +797,7 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
                     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",
+                        "title": f"{prediction_bbox}; baseline 0 0; x_size 0; x_descenders 0; x_ascenders 0",
                     },
                 )
                 word_div = SubElement(
@@ -781,9 +806,7 @@ def export_as_xml(self, file_title: str = "docTR - XML export (hOCR)") -> tuple[
                     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))}",
+                        "title": f"{prediction_bbox}; x_wconf {int(round(prediction.confidence * 100))}",
                     },
                 )
                 word_div.text = prediction.value

doctr/models/_utils.py

@@ -59,7 +59,7 @@ def estimate_orientation(
 
     # Convert image to grayscale if necessary
     if img.shape[-1] == 3:
-        gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        gray_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
         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]
     else:

doctr/models/detection/core.py

@@ -7,6 +7,7 @@
 import cv2
 import numpy as np
 
+from doctr.utils.multithreading import multithread_exec
 from doctr.utils.repr import NestedObject
 
 __all__ = ["DetectionPostProcessor"]
@@ -69,6 +70,28 @@ def bitmap_to_boxes(
     ) -> np.ndarray:
         raise NotImplementedError
 
+    def _process_sample(self, proba_map: np.ndarray) -> list[np.ndarray]:
+        """Performs postprocessing for a single sample
+
+        Args:
+            proba_map: probability map of shape (H, W, C)
+
+        Returns:
+            list of C class predictions, each of shape (*, 5) or (*, 6)
+        """
+        return [
+            self.bitmap_to_boxes(
+                proba_map[..., idx],
+                # Erosion + dilation on the binary map
+                cv2.morphologyEx(
+                    (proba_map[..., idx] >= self.bin_thresh).astype(np.uint8),
+                    cv2.MORPH_OPEN,
+                    self._opening_kernel,
+                ),
+            )
+            for idx in range(proba_map.shape[-1])
+        ]
+
     def __call__(
         self,
         proba_map,
@@ -84,17 +107,4 @@ def __call__(
         """
         if proba_map.ndim != 4:
             raise AssertionError(f"arg `proba_map` is expected to be 4-dimensional, got {proba_map.ndim}.")
-
-        # Erosion + dilation on the binary map
-        bin_map = [
-            [
-                cv2.morphologyEx(bmap[..., idx], cv2.MORPH_OPEN, self._opening_kernel)
-                for idx in range(proba_map.shape[-1])
-            ]
-            for bmap in (proba_map >= self.bin_thresh).astype(np.uint8)
-        ]
-
-        return [
-            [self.bitmap_to_boxes(pmaps[..., idx], bmaps[idx]) for idx in range(proba_map.shape[-1])]
-            for pmaps, bmaps in zip(proba_map, bin_map)
-        ]
+        return list(multithread_exec(self._process_sample, proba_map))

doctr/utils/metrics.py

@@ -7,7 +7,7 @@
 import numpy as np
 from anyascii import anyascii
 from scipy.optimize import linear_sum_assignment
-from shapely import area, intersection, polygons
+from shapely import STRtree, area, intersection, polygons
 
 __all__ = [
     "TextMatch",
@@ -146,7 +146,7 @@ def box_iou(boxes_1: np.ndarray, boxes_2: np.ndarray) -> np.ndarray:
 
         intersection = np.clip(right - left, 0, np.inf) * np.clip(bot - top, 0, np.inf)
         union = (r1 - l1) * (b1 - t1) + ((r2 - l2) * (b2 - t2)).T - intersection
-        iou_mat = intersection / union
+        iou_mat = np.divide(intersection, union, out=np.zeros_like(intersection), where=union > 0)
 
     return iou_mat
 
@@ -170,18 +170,15 @@ def polygon_iou(polys_1: np.ndarray, polys_2: np.ndarray) -> np.ndarray:
 
     geoms_1 = polygons(polys_1)
     geoms_2 = polygons(polys_2)
-    grid_1 = np.repeat(geoms_1, m)
-    grid_2 = np.tile(geoms_2, n)
 
-    # Compute intersections and areas
-    intersections = area(intersection(grid_1, grid_2))
-    areas_1 = area(grid_1)
-    areas_2 = area(grid_2)
-
-    # Compute IoU
-    unions = areas_1 + areas_2 - intersections
-    iou_flat = np.divide(intersections, unions, out=np.zeros_like(intersections), where=unions > 0)
-    return iou_flat.reshape(n, m).astype(np.float32)
+    iou_mat = np.zeros((n, m), dtype=np.float32)
+    idcs_1, idcs_2 = STRtree(geoms_2).query(geoms_1)
+    if idcs_1.size > 0:
+        # Compute intersections and areas for the candidate pairs only
+        intersections = area(intersection(geoms_1[idcs_1], geoms_2[idcs_2]))
+        unions = area(geoms_1)[idcs_1] + area(geoms_2)[idcs_2] - intersections
+        iou_mat[idcs_1, idcs_2] = np.divide(intersections, unions, out=np.zeros_like(intersections), where=unions > 0)
+    return iou_mat
 
 
 def nms(boxes: np.ndarray, thresh: float = 0.5) -> list[int]:
@@ -275,7 +272,7 @@ def update(self, gts: np.ndarray, preds: np.ndarray) -> None:
             gts: a set of relative bounding boxes either of shape (N, 4) or (N, 5) if they are rotated ones
             preds: a set of relative bounding boxes either of shape (M, 4) or (M, 5) if they are rotated ones
         """
-        if preds.shape[0] > 0:
+        if gts.shape[0] > 0 and preds.shape[0] > 0:
             # Compute IoU
             if self.use_polygons:
                 iou_mat = polygon_iou(gts, preds)
@@ -468,8 +465,7 @@ def update(
                 "there should be the same number of boxes and string both for the ground truth and the predictions"
             )
 
-        # Compute IoU
-        if pred_boxes.shape[0] > 0:
+        if gt_boxes.shape[0] > 0 and pred_boxes.shape[0] > 0:
             if self.use_polygons:
                 iou_mat = polygon_iou(gt_boxes, pred_boxes)
             else:
@@ -598,8 +594,7 @@ def update(
                 "there should be the same number of boxes and string both for the ground truth and the predictions"
             )
 
-        # Compute IoU
-        if pred_boxes.shape[0] > 0:
+        if gt_boxes.shape[0] > 0 and pred_boxes.shape[0] > 0:
             if self.use_polygons:
                 iou_mat = polygon_iou(gt_boxes, pred_boxes)
             else:

tests/common/test_io_elements.py

@@ -467,11 +467,17 @@ def test_page():
     }
 
     # Export XML
-    assert (
-        isinstance(page.export_as_xml(), tuple)
-        and isinstance(page.export_as_xml()[0], (bytes, bytearray))
-        and isinstance(page.export_as_xml()[1], ElementTree)
-    )
+    xml_bytes, xml_tree = page.export_as_xml()
+    assert isinstance(xml_bytes, (bytes, bytearray)) and isinstance(xml_tree, ElementTree)
+    # The detected language must be exported instead of being hardcoded to "en"
+    assert xml_tree.getroot().get("xml:lang") == "EN"
+    # hOCR title properties must be single-spaced
+    titles = [el.get("title") for el in xml_tree.iter() if el.get("title") is not None]
+    assert any(title.startswith("bbox ") for title in titles)
+    assert all("  " not in title for title in titles)
+    # Without a detected language, the export must fall back to "en"
+    fallback_page = elements.Page(np.zeros((300, 200, 3), dtype=np.uint8), blocks, page_idx, page_size, orientation)
+    assert fallback_page.export_as_xml()[1].getroot().get("xml:lang") == "en"
 
     # Repr
     assert "\n".join(repr(page).split("\n")[:2]) == f"Page(\n  dimensions={page_size!r}"
@@ -529,11 +535,19 @@ def test_kiepage():
     }
 
     # Export XML
-    assert (
-        isinstance(kie_page.export_as_xml(), tuple)
-        and isinstance(kie_page.export_as_xml()[0], (bytes, bytearray))
-        and isinstance(kie_page.export_as_xml()[1], ElementTree)
+    xml_bytes, xml_tree = kie_page.export_as_xml()
+    assert isinstance(xml_bytes, (bytes, bytearray)) and isinstance(xml_tree, ElementTree)
+    # The detected language must be exported instead of being hardcoded to "en"
+    assert xml_tree.getroot().get("xml:lang") == "EN"
+    # hOCR title properties must be single-spaced
+    titles = [el.get("title") for el in xml_tree.iter() if el.get("title") is not None]
+    assert any(title.startswith("bbox ") for title in titles)
+    assert all("  " not in title for title in titles)
+    # Without a detected language, the export must fall back to "en"
+    fallback_page = elements.KIEPage(
+        np.zeros((300, 200, 3), dtype=np.uint8), predictions, page_idx, page_size, orientation
     )
+    assert fallback_page.export_as_xml()[1].getroot().get("xml:lang") == "en"
 
     # Repr
     assert "\n".join(repr(kie_page).split("\n")[:2]) == f"KIEPage(\n  dimensions={page_size!r}"