feat: builder modifications

mindee · charlesmindee · Dec 26, 2021 · Dec 16, 2021 · Dec 17, 2021 · Dec 17, 2021
commit dc95b3314de44a0072410189dcec556373938e3b
diff --git a/doctr/models/builder.py b/doctr/models/builder.py
@@ -10,7 +10,7 @@
 from scipy.cluster.hierarchy import fclusterdata
 
 from doctr.io.elements import Block, Document, Line, Page, Word
-from doctr.utils.geometry import rbbox_to_polygon, resolve_enclosing_bbox, resolve_enclosing_rbbox, rotate_boxes
+from doctr.utils.geometry import resolve_enclosing_bbox, resolve_enclosing_rbbox, rotate_boxes
 from doctr.utils.repr import NestedObject
 
 __all__ = ['DocumentBuilder']
@@ -57,7 +57,8 @@ def _sort_boxes(boxes: np.ndarray) -> np.ndarray:
                 orig_shape=(1024, 1024),
                 min_angle=5.,
             )
-            return (boxes[:, 0] + 2 * boxes[:, 1] / np.median(boxes[:, 3])).argsort()
+            # Points are in this order: top left, top right, bot right, bot left
+            return (boxes[:, 0, 0] + 2 * boxes[:, 1, 2] / np.median(boxes[:, 1, 2] - boxes[:, 1, 1])).argsort()
         return (boxes[:, 0] + 2 * boxes[:, 3] / np.median(boxes[:, 3] - boxes[:, 1])).argsort()
 
     def _resolve_sub_lines(self, boxes: np.ndarray, words: List[int]) -> List[List[int]]:
@@ -84,7 +85,7 @@ def _resolve_sub_lines(self, boxes: np.ndarray, words: List[int]) -> List[List[i
                 prev_box = boxes[sub_line[-1]]
                 # Compute distance between boxes
                 if len(boxes.shape) == 3:
-                    dist = boxes[i, 0] - prev_box[2] / 2 - (prev_box[0] + prev_box[2] / 2)
+                    dist = boxes[i, 0, 0] - prev_box[0, 1]
                 else:
                     dist = boxes[i, 0] - prev_box[2]
                 # If distance between boxes is lower than paragraph break, same sub-line
@@ -110,7 +111,7 @@ def _resolve_lines(self, boxes: np.ndarray) -> List[List[int]]:
             nested list of box indices
         """
         # Compute median for boxes heights
-        y_med = np.median(boxes[:, 3] if len(boxes.shape) == 3 else boxes[:, 3] - boxes[:, 1])
+        y_med = np.median(boxes[:, 1, 2] - boxes[:, 1, 1] if len(boxes.shape) == 3 else boxes[:, 3] - boxes[:, 1])
 
         # Sort boxes
         idxs = self._sort_boxes(boxes)
@@ -119,7 +120,7 @@ def _resolve_lines(self, boxes: np.ndarray) -> List[List[int]]:
         words = [idxs[0]]  # Assign the top-left word to the first line
         # Define a mean y-center for the line
         if len(boxes.shape) == 3:
-            y_center_sum = boxes[idxs[0]][1]
+            y_center_sum = boxes[idxs[0]][[[1, 2], [1, 1]]].mean()
         else:
             y_center_sum = boxes[idxs[0]][[1, 3]].mean()
 
@@ -128,7 +129,7 @@ def _resolve_lines(self, boxes: np.ndarray) -> List[List[int]]:
 
             # Compute y_dist
             if len(boxes.shape) == 3:
-                y_dist = abs(boxes[idx][1] - y_center_sum / len(words))
+                y_dist = abs(boxes[idx][[[1, 2], [1, 1]]].mean() - y_center_sum / len(words))
             else:
                 y_dist = abs(boxes[idx][[1, 3]].mean() - y_center_sum / len(words))
             # If y-center of the box is close enough to mean y-center of the line, same line
@@ -142,7 +143,7 @@ def _resolve_lines(self, boxes: np.ndarray) -> List[List[int]]:
                 y_center_sum = 0
 
             words.append(idx)
-            y_center_sum += boxes[idx][1 if len(boxes.shape) == 3 else [1, 3]].mean()
+            y_center_sum += boxes[idx][[[1, 2], [1, 1]] if len(boxes.shape) == 3 else [1, 3]].mean()
 
         # Use the remaining words to form the last(s) line(s)
         if len(words) > 0:
@@ -241,7 +242,7 @@ def _build_blocks(self, boxes: np.ndarray, word_preds: List[Tuple[str, float]])
                     [
                         Word(
                             *word_preds[idx],
-                            (boxes[idx, 0], boxes[idx, 1], boxes[idx, 2], boxes[idx, 3], boxes[idx, 4])
+                            tuple(boxes[idx].tolist())
                         ) if len(boxes.shape) == 3 else
                         Word(
                             *word_preds[idx],

diff --git a/doctr/utils/geometry.py b/doctr/utils/geometry.py
@@ -20,11 +20,13 @@ def bbox_to_polygon(bbox: BoundingBox) -> Polygon4P:
     return bbox[0], (bbox[1][0], bbox[0][1]), (bbox[0][0], bbox[1][1]), bbox[1]
 
 
+# TODO: deprecation warning (a rbbox is now a polygon)
 def rbbox_to_polygon(rbbox: RotatedBbox) -> Polygon4P:
     (x, y, w, h, alpha) = rbbox
     return cv2.boxPoints(((float(x), float(y)), (float(w), float(h)), -float(alpha)))
 
 
+# TODO: deprecation warning (a rbbox is now a polygon)
 def fit_rbbox(pts: np.ndarray) -> RotatedBbox:
     ((x, y), (w, h), alpha) = cv2.minAreaRect(pts)
     return x, y, h, w, 90 - alpha
@@ -35,6 +37,7 @@ def polygon_to_bbox(polygon: Polygon4P) -> BoundingBox:
     return (min(x), min(y)), (max(x), max(y))
 
 
+# TODO: deprecation warning (a rbbox is now a polygon)
 def polygon_to_rbbox(polygon: Polygon4P) -> RotatedBbox:
     cnt = np.array(polygon).reshape((-1, 1, 2)).astype(np.float32)
     return fit_rbbox(cnt)
@@ -58,9 +61,10 @@ def resolve_enclosing_bbox(bboxes: Union[List[BoundingBox], np.ndarray]) -> Unio
         return (min(x), min(y)), (max(x), max(y))
 
 
-def resolve_enclosing_rbbox(rbboxes: List[RotatedBbox]) -> RotatedBbox:
-    pts = np.asarray([pt for rbbox in rbboxes for pt in rbbox_to_polygon(rbbox)], np.float32)
-    return fit_rbbox(pts)
+def resolve_enclosing_rbbox(rbboxes: List[RotatedBbox]) -> Polygon4P:
+    cloud = np.concatenate(rbboxes, axis=0)
+    rect = cv2.minAreaRect(cloud)
+    return cv2.boxPoints(rect)
 
 
 def rotate_abs_points(points: np.ndarray, angle: float = 0.) -> np.ndarray: