fix(ocr-utils): unit test and fix the rotate_bounding_box function (#1897)

Signed-off-by: Clément Doumouro <[email protected]>

Author: ClemDoum

docling/utils/ocr_utils.py

@@ -41,7 +41,7 @@ def tesseract_box_to_bounding_rectangle(
     im_size: Tuple[int, int],
 ) -> BoundingRectangle:
     # box is in the top, left, height, width format, top left coordinates
-    rect = rotate_bounding_box(bbox, angle=-orientation, im_size=im_size)
+    rect = rotate_bounding_box(bbox, angle=orientation, im_size=im_size)
     rect = BoundingRectangle(
         r_x0=rect.r_x0 / scale,
         r_y0=rect.r_y0 / scale,

docling/utils/orientation.py

@@ -14,51 +14,44 @@ def rotate_bounding_box(
     # coordinate system. Then other corners are found rotating counterclockwise
     bbox = bbox.to_top_left_origin(im_size[1])
     left, top, width, height = bbox.l, bbox.t, bbox.width, bbox.height
-    im_h, im_w = im_size
+    im_w, im_h = im_size
     angle = angle % 360
     if angle == 0:
-        r_x0 = left
-        r_y0 = top + height
-        r_x1 = r_x0 + width
-        r_y1 = r_y0
-        r_x2 = r_x0 + width
-        r_y2 = r_y0 - height
-        r_x3 = r_x0
-        r_y3 = r_y0 - height
+        return BoundingRectangle.from_bounding_box(bbox)
     elif angle == 90:
-        r_x0 = im_w - (top + height)
-        r_y0 = left
+        r_x0 = top + height
+        r_y0 = im_w - left
         r_x1 = r_x0
-        r_y1 = r_y0 + width
-        r_x2 = r_x0 + height
-        r_y2 = r_y0 + width
-        r_x3 = r_x0
-        r_y3 = r_y0 + width
+        r_y1 = r_y0 - width
+        r_x2 = r_x1 - height
+        r_y2 = r_y1
+        r_x3 = r_x2
+        r_y3 = r_y0
     elif angle == 180:
-        r_x0 = im_h - left
-        r_y0 = im_w - (top + height)
+        r_x0 = im_w - left
+        r_y0 = im_h - (top + height)
         r_x1 = r_x0 - width
         r_y1 = r_y0
-        r_x2 = r_x0 - width
-        r_y2 = r_y0 + height
+        r_x2 = r_x1
+        r_y2 = r_y1 + height
         r_x3 = r_x0
-        r_y3 = r_y0 + height
+        r_y3 = r_y2
     elif angle == 270:
-        r_x0 = top + height
-        r_y0 = im_h - left
+        r_x0 = im_h - (top + height)
+        r_y0 = left
         r_x1 = r_x0
-        r_y1 = r_y0 - width
-        r_x2 = r_x0 - height
-        r_y2 = r_y0 - width
-        r_x3 = r_x0 - height
+        r_y1 = r_y0 + width
+        r_x2 = r_x1 + height
+        r_y2 = r_y1
+        r_x3 = r_x2
         r_y3 = r_y0
     else:
         msg = (
             f"invalid orientation {angle}, expected values in:"
             f" {sorted(CLIPPED_ORIENTATIONS)}"
         )
         raise ValueError(msg)
-    return BoundingRectangle(
+    rectangle = BoundingRectangle(
         r_x0=r_x0,
         r_y0=r_y0,
         r_x1=r_x1,
@@ -69,3 +62,4 @@ def rotate_bounding_box(
         r_y3=r_y3,
         coord_origin=CoordOrigin.TOPLEFT,
     )
+    return rectangle

pyproject.toml

@@ -44,7 +44,7 @@ authors = [
 requires-python = '>=3.9,<4.0'
 dependencies = [
   'pydantic (>=2.0.0,<3.0.0)',
-  'docling-core[chunking] (>=2.39.0,<3.0.0)',
+  'docling-core[chunking] (>=2.40.0,<3.0.0)',
   'docling-parse (>=4.0.0,<5.0.0)',
   'docling-ibm-models (>=3.6.0,<4)',
   'filetype (>=1.2.0,<2.0.0)',

tests/test_ocr_utils.py

@@ -0,0 +1,80 @@
+from typing import Tuple
+
+import pytest
+from docling_core.types.doc import BoundingBox, CoordOrigin
+from docling_core.types.doc.page import BoundingRectangle
+
+from docling.utils.orientation import rotate_bounding_box
+
+IM_SIZE = (4, 5)
+BBOX = BoundingBox(l=1, t=3, r=3, b=4, coord_origin=CoordOrigin.TOPLEFT)
+RECT = BoundingRectangle(
+    r_x0=1,
+    r_y0=4,
+    r_x1=3,
+    r_y1=4,
+    r_x2=3,
+    r_y2=3,
+    r_x3=1,
+    r_y3=3,
+    coord_origin=CoordOrigin.TOPLEFT,
+)
+RECT_90 = BoundingRectangle(
+    r_x0=4,
+    r_y0=3,
+    r_x1=4,
+    r_y1=1,
+    r_x2=3,
+    r_y2=1,
+    r_x3=3,
+    r_y3=3,
+    coord_origin=CoordOrigin.TOPLEFT,
+)
+RECT_180 = BoundingRectangle(
+    r_x0=3,
+    r_y0=1,
+    r_x1=1,
+    r_y1=1,
+    r_x2=1,
+    r_y2=2,
+    r_x3=3,
+    r_y3=2,
+    coord_origin=CoordOrigin.TOPLEFT,
+)
+RECT_270 = BoundingRectangle(
+    r_x0=1,
+    r_y0=1,
+    r_x1=1,
+    r_y1=3,
+    r_x2=2,
+    r_y2=3,
+    r_x3=2,
+    r_y3=1,
+    coord_origin=CoordOrigin.TOPLEFT,
+)
+
+
+@pytest.mark.parametrize(
+    ["bbox", "im_size", "angle", "expected_rectangle"],
+    [
+        # (BBOX, IM_SIZE, 0, RECT),
+        # (BBOX, IM_SIZE, 90, RECT_90),
+        (BBOX, IM_SIZE, 180, RECT_180),
+        # (BBOX, IM_SIZE, 270, RECT_270),
+        # (BBOX, IM_SIZE, 360, RECT),
+        # (BBOX, IM_SIZE, -90, RECT_270),
+        (BBOX, IM_SIZE, -180, RECT_180),
+        # (BBOX, IM_SIZE, -270, RECT_90),
+    ],
+)
+def test_rotate_bounding_box(
+    bbox: BoundingBox,
+    im_size: Tuple[int, int],
+    angle: int,
+    expected_rectangle: BoundingRectangle,
+):
+    rotated = rotate_bounding_box(bbox, angle, im_size)
+
+    assert rotated == expected_rectangle
+    expected_angle_360 = angle % 360
+    assert rotated.angle_360 == expected_angle_360