Repository: mindee/doctr

Issue:
[models] Add full OCR module
Design an object that will wrap all DL model components and be responsible for localizing and identifying all text elements in documents.

**Inputs**
- a collection of documents, where each document is a list of pages, themselves expressed as numpy-encoded images.

**Outputs**
- a collection of `Document` objects

The following components will be required:
- [x] DetectionPredictor (#39)
- [x] RecognitionPredictor (#39)
- [x] OCRPredictor (#39)
- [x] ElementExporter (#16, #26, #40)


Discussion:

Resolving patch:
diff --git a/doctr/documents/elements.py b/doctr/documents/elements.py
--- a/doctr/documents/elements.py
+++ b/doctr/documents/elements.py
@@ -165,14 +165,14 @@ def __init__(
         blocks: List[Block],
         page_idx: int,
         dimensions: Tuple[int, int],
-        orientation: Dict[str, Any],
-        language: Dict[str, Any],
+        orientation: Optional[Dict[str, Any]] = None,
+        language: Optional[Dict[str, Any]] = None,
     ) -> None:
         super().__init__(blocks=blocks)
         self.page_idx = page_idx
         self.dimensions = dimensions
-        self.orientation = orientation
-        self.language = language
+        self.orientation = orientation if isinstance(orientation, dict) else dict(value=None, confidence=None)
+        self.language = language if isinstance(language, dict) else dict(value=None, confidence=None)
 
     def render(self, block_break: str = '\n\n') -> str:
         """Renders the full text of the element"""
diff --git a/doctr/models/core.py b/doctr/models/core.py
--- a/doctr/models/core.py
+++ b/doctr/models/core.py
@@ -11,6 +11,7 @@
 from .detection import DetectionPredictor
 from .recognition import RecognitionPredictor
 from ._utils import extract_crops
+from ..documents.elements import Word, Line, Block, Page, Document
 
 __all__ = ['OCRPredictor']
 
@@ -32,10 +33,61 @@ def __init__(
         self.det_predictor = det_predictor
         self.reco_predictor = reco_predictor
 
+    @staticmethod
+    def _build_blocks(boxes: np.ndarray, char_sequences: List[str]) -> List[Block]:
+        """Gather independent words in structured blocks
+
+        Args:
+            boxes: bounding boxes of all detected words of the page, of shape (N, 4)
+            char_sequences: list of all detected words of the page, of shape N
+
+        Returns:
+            list of block elements
+        """
+
+        if boxes.shape[0] != len(char_sequences):
+            raise ValueError(f"Incompatible argument lengths: {boxes.shape[0]}, {len(char_sequences)}")
+
+        # Sort bounding boxes from top to bottom, left to right
+        idxs = (boxes[:, 0] + boxes[:, 1] * (boxes[:, 2] - boxes[:, 0])).argsort()
+        # Try to arrange them in lines
+        lines = []
+        words: List[int] = []
+        for idx in idxs:
+
+            if len(words) == 0:
+                words = [idx]
+                continue
+            # Check horizontal gaps
+            horz_gap = boxes[idx, 0] - boxes[words[-1], 0] < 0.5 * (boxes[words[-1], 2] - boxes[words[-1], 0])
+            # Check vertical gaps
+            vert_gap = abs(boxes[idx, 3] - boxes[words[-1], 3]) > 0.5 * (boxes[words[-1], 3] - boxes[words[-1], 1])
+
+            if horz_gap or vert_gap:
+                lines.append(words)
+                words.clear()
+
+            words.append(idx)
+
+        if len(words) > 0:
+            lines.append(words)
+
+        blocks = [
+            Block([Line(
+                [Word(
+                    char_sequences[idx],
+                    boxes[idx, 4],
+                    ((boxes[idx, 0], boxes[idx, 1]), (boxes[idx, 2], boxes[idx, 3]))
+                ) for idx in line]
+            )]) for line in lines
+        ]
+
+        return blocks
+
     def __call__(
         self,
         documents: List[List[np.ndarray]],
-    ) -> List[List[List[Dict[str, Any]]]]:
+    ) -> List[Document]:
 
         pages = [page for doc in documents for page in doc]
 
@@ -52,15 +104,21 @@ def __call__(
         page_idx, crop_idx = 0, 0
         results = []
         for doc_idx, nb_pages in enumerate(num_pages):
-            doc = []
+            _pages = []
             for page_boxes in boxes[page_idx: page_idx + nb_pages]:
-                page = []
-                for _idx in range(num_crops[page_idx]):
-                    page.append(dict(box=page_boxes[_idx], text=char_sequences[crop_idx + _idx]))
-
+                # Assemble all detected words into structured blocks
+                _pages.append(
+                    Page(
+                        self._build_blocks(
+                            page_boxes[:num_crops[page_idx]],
+                            char_sequences[crop_idx: crop_idx + num_crops[page_idx]]
+                        ),
+                        page_idx,
+                        pages[page_idx].shape[:2],
+                    )
+                )
                 crop_idx += num_crops[page_idx]
                 page_idx += 1
-                doc.append(page)
-            results.append(doc)
+            results.append(Document(_pages))
 
         return results


Test patch:
diff --git a/test/test_models.py b/test/test_models.py
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -15,7 +15,7 @@
 
 from tensorflow.keras import layers, Sequential
 
-from doctr.documents import read_pdf
+from doctr.documents import read_pdf, Document, Page, Block, Line, Word
 from test_documents import mock_pdf
 from doctr import models
 
@@ -262,11 +262,10 @@ def test_ocrpredictor(mock_pdf, mock_mapping, test_detectionpredictor, test_reco
     out = predictor(docs)
 
     assert len(out) == num_docs
+    # Document
+    assert all(isinstance(doc, Document) for doc in out)
     # The input PDF has 8 pages
-    assert all(len(doc) == 8 for doc in out)
-    # Structure of page
-    assert all(isinstance(page, list) for doc in out for page in doc)
-    assert all(isinstance(elt, dict) for doc in out for page in doc for elt in page)
+    assert all(len(doc.pages) == 8 for doc in out)
 
 
 def test_sar_decoder(mock_mapping):