Repository: mindee/doctr

Issue:
[documents] improve line detection
An OCR predictor must detect lines, and our current version is too weak:
- many lines overlap
- some lines stop in the middle of a dense block whereas other lines are kind of "bridging" between separated blocks.

![line11](https://user-images.githubusercontent.com/70526046/116210562-465bc980-a743-11eb-8806-80afa34863dd.png)



Discussion:

Resolving patch:
diff --git a/doctr/models/core.py b/doctr/models/core.py
--- a/doctr/models/core.py
+++ b/doctr/models/core.py
@@ -71,7 +71,7 @@ def __init__(
         self,
         resolve_lines: bool = False,
         resolve_blocks: bool = False,
-        paragraph_break: float = 0.15
+        paragraph_break: float = 0.035
     ) -> None:
 
         self.resolve_lines = resolve_lines
@@ -91,43 +91,82 @@ def _sort_boxes(boxes: np.ndarray) -> np.ndarray:
         Returns:
             indices of ordered boxes of shape (N,)
         """
+        return (boxes[:, 0] + 2 * boxes[:, 3] / np.median(boxes[:, 3] - boxes[:, 1])).argsort()
 
-        return (boxes[:, 0] + boxes[:, 3] / np.median(boxes[:, 3] - boxes[:, 1])).argsort()
+    def _resolve_sub_lines(self, boxes: np.ndarray, words: List[int]) -> List[List[int]]:
+        """Split a line in sub_lines
 
-    def _resolve_lines(self, boxes: np.ndarray, idxs: np.ndarray) -> List[List[int]]:
-        """Uses ordered boxes to group them in lines
+        Args:
+            boxes: bounding boxes of shape (N, 4)
+            words: list of indexes for the words of the line
+
+        Returns:
+            A list of (sub-)lines computed from the original line (words)
+        """
+        lines = []
+        # Sort words horizontally
+        words = [words[j] for j in np.argsort([boxes[i, 0] for i in words]).tolist()]
+        # Eventually split line horizontally
+        if len(words) < 2:
+            lines.append(words)
+        else:
+            sub_line = [words[0]]
+            for i in words[1:]:
+                horiz_break = True
+
+                prev_box = boxes[sub_line[-1]]
+                # If distance between boxes is lower than paragraph break, same sub-line
+                if (boxes[i, 0] - prev_box[2]) < self.paragraph_break:
+                    horiz_break = False
+
+                if horiz_break:
+                    lines.append(sub_line)
+                    sub_line = []
+
+                sub_line.append(i)
+            lines.append(sub_line)
+
+        return lines
+
+    def _resolve_lines(self, boxes: np.ndarray) -> List[List[int]]:
+        """Order boxes to group them in lines
 
         Args:
             boxes: bounding boxes of shape (N, 4)
-            idxs: indices of ordered boxes of shape (N,)
 
         Returns:
             nested list of box indices
         """
+        # Compute median for boxes heights
+        y_med = np.median(boxes[:, 3] - boxes[:, 1])
+        # Sort boxes
+        idxs = (boxes[:, 0] + 2 * boxes[:, 3] / y_med).argsort()
 
-        # Try to arrange them in lines
         lines = []
-        # Add the first word anyway
-        words: List[int] = [idxs[0]]
+        words = [idxs[0]]  # Assign the top-left word to the first line
+        # Define a mean y-center for the line
+        y_center_sum = boxes[idxs[0]][[1, 3]].mean()
+
         for idx in idxs[1:]:
-            line_break = True
+            vert_break = True
 
-            prev_box = boxes[words[-1]]
-            # Reduced vertical diff
-            if boxes[idx, 1] < prev_box[[1, 3]].mean():
-                # Words horizontally ordered and close
-                if (boxes[idx, 0] - prev_box[2]) < self.paragraph_break:
-                    line_break = False
+            # If y-center of the box is close enough to mean y-center of the line, same line
+            if abs(boxes[idx][[1, 3]].mean() - y_center_sum / len(words)) < y_med / 2:
+                vert_break = False
 
-            if line_break:
-                lines.append(words)
+            if vert_break:
+                # Compute sub-lines (horizontal split)
+                lines.extend(self._resolve_sub_lines(boxes, words))
                 words = []
+                y_center_sum = 0
 
             words.append(idx)
+            y_center_sum += boxes[idx][[1, 3]].mean()
 
-        # Use the remaining words to form the last line
+        # Use the remaining words to form the last(s) line(s)
         if len(words) > 0:
-            lines.append(words)
+            # Compute sub-lines (horizontal split)
+            lines.extend(self._resolve_sub_lines(boxes, words))
 
         return lines
 
@@ -148,15 +187,12 @@ def _build_blocks(self, boxes: np.ndarray, char_sequences: List[str]) -> List[Bl
         if boxes.shape[0] == 0:
             return []
 
-        # Sort bounding boxes from top to bottom, left to right
-        idxs = self._sort_boxes(boxes[:, :4])
-
         # Decide whether we try to form lines
         if self.resolve_lines:
-            lines = self._resolve_lines(boxes[:, :4], idxs)
+            lines = self._resolve_lines(boxes[:, :4])
         else:
-            # One line for all words
-            lines = [idxs]
+            # Sort bounding boxes, one line for all boxes
+            lines = [self._sort_boxes(boxes[:, :4])]
 
         # No automatic line grouping yet --> 1 block for all lines
         blocks = [


Test patch:
diff --git a/test/test_models.py b/test/test_models.py
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -67,7 +67,7 @@ def test_documentbuilder():
     assert len(out.pages[0].blocks) == 0
 
     # Repr
-    assert repr(doc_builder) == "DocumentBuilder(resolve_lines=True, paragraph_break=0.15)"
+    assert repr(doc_builder) == "DocumentBuilder(resolve_lines=True, paragraph_break=0.035)"
 
 
 @pytest.mark.parametrize(
@@ -76,7 +76,7 @@ def test_documentbuilder():
         [[[0, 0.5, 0.1, 0.6], [0, 0.3, 0.2, 0.4], [0, 0, 0.1, 0.1]], [2, 1, 0]],  # vertical
         [[[0.7, 0.5, 0.85, 0.6], [0.2, 0.3, 0.4, 0.4], [0, 0, 0.1, 0.1]], [2, 1, 0]],  # diagonal
         [[[0, 0.5, 0.1, 0.6], [0.15, 0.5, 0.25, 0.6], [0.5, 0.5, 0.6, 0.6]], [0, 1, 2]],  # same line, 2p
-        [[[0, 0.5, 0.1, 0.6], [0.2, 0.48, 0.35, 0.58], [0.8, 0.52, 0.9, 0.63]], [0, 1, 2]],  # ~same line
+        [[[0, 0.5, 0.1, 0.6], [0.2, 0.49, 0.35, 0.59], [0.8, 0.52, 0.9, 0.63]], [0, 1, 2]],  # ~same line
         [[[0, 0.3, 0.4, 0.45], [0.5, 0.28, 0.75, 0.42], [0, 0.45, 0.1, 0.55]], [0, 1, 2]],  # 2 lines
         [[[0, 0.3, 0.4, 0.35], [0.75, 0.28, 0.95, 0.42], [0, 0.45, 0.1, 0.55]], [0, 1, 2]],  # 2 lines
     ],
@@ -88,20 +88,20 @@ def test_sort_boxes(input_boxes, sorted_idxs):
 
 
 @pytest.mark.parametrize(
-    "input_boxes, sorted_idxs, lines",
+    "input_boxes, lines",
     [
-        [[[0, 0.5, 0.1, 0.6], [0, 0.3, 0.2, 0.4], [0, 0, 0.1, 0.1]], [2, 1, 0], [[2], [1], [0]]],  # vertical
-        [[[0.7, 0.5, 0.85, 0.6], [0.2, 0.3, 0.4, 0.4], [0, 0, 0.1, 0.1]], [2, 1, 0], [[2], [1], [0]]],  # diagonal
-        [[[0, 0.5, 0.1, 0.6], [0.15, 0.5, 0.25, 0.6], [0.5, 0.5, 0.6, 0.6]], [0, 1, 2], [[0, 1], [2]]],  # same line, 2p
-        [[[0, 0.5, 0.1, 0.6], [0.2, 0.48, 0.35, 0.58], [0.8, 0.52, 0.9, 0.63]], [0, 1, 2], [[0, 1], [2]]],  # ~same line
-        [[[0, 0.3, 0.4, 0.45], [0.5, 0.28, 0.75, 0.42], [0, 0.45, 0.1, 0.55]], [0, 1, 2], [[0, 1], [2]]],  # 2 lines
-        [[[0, 0.3, 0.4, 0.35], [0.75, 0.28, 0.95, 0.42], [0, 0.45, 0.1, 0.55]], [0, 1, 2], [[0], [1], [2]]],  # 2 lines
+        [[[0, 0.5, 0.1, 0.6], [0, 0.3, 0.2, 0.4], [0, 0, 0.1, 0.1]], [[2], [1], [0]]],  # vertical
+        [[[0.7, 0.5, 0.85, 0.6], [0.2, 0.3, 0.4, 0.4], [0, 0, 0.1, 0.1]], [[2], [1], [0]]],  # diagonal
+        [[[0, 0.5, 0.14, 0.6], [0.15, 0.5, 0.25, 0.6], [0.5, 0.5, 0.6, 0.6]], [[0, 1], [2]]],  # same line, 2p
+        [[[0, 0.5, 0.18, 0.6], [0.2, 0.48, 0.35, 0.58], [0.8, 0.52, 0.9, 0.63]], [[0, 1], [2]]],  # ~same line
+        [[[0, 0.3, 0.48, 0.45], [0.5, 0.28, 0.75, 0.42], [0, 0.45, 0.1, 0.55]], [[0, 1], [2]]],  # 2 lines
+        [[[0, 0.3, 0.4, 0.35], [0.75, 0.28, 0.95, 0.42], [0, 0.45, 0.1, 0.55]], [[0], [1], [2]]],  # 2 lines
     ],
 )
-def test_resolve_lines(input_boxes, sorted_idxs, lines):
+def test_resolve_lines(input_boxes, lines):
 
     doc_builder = models.DocumentBuilder()
-    assert doc_builder._resolve_lines(np.asarray(input_boxes), np.asarray(sorted_idxs)) == lines
+    assert doc_builder._resolve_lines(np.asarray(input_boxes)) == lines
 
 
 def test_ocrpredictor(mock_pdf, test_detectionpredictor, test_recognitionpredictor):  # noqa: F811