Initial commit

docling/models/easyocr_model.py

@@ -0,0 +1,77 @@
+import copy
+import logging
+import random
+from typing import Iterable
+
+import numpy
+from PIL import ImageDraw
+
+from docling.datamodel.base_models import BoundingBox, CoordOrigin, OcrCell, Page
+
+_log = logging.getLogger(__name__)
+
+
+class EasyOcrModel:
+    def __init__(self, config):
+        self.config = config
+        self.enabled = config["enabled"]
+        self.scale = 3  # multiplier for 72 dpi == 216 dpi.
+
+        if self.enabled:
+            import easyocr
+
+            self.reader = easyocr.Reader(config["lang"])
+
+    def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
+
+        if not self.enabled:
+            yield from page_batch
+            return
+
+        for page in page_batch:
+            # rects = page._fpage.
+            high_res_image = page._backend.get_page_image(scale=self.scale)
+            im = numpy.array(high_res_image)
+            result = self.reader.readtext(im)
+
+            del high_res_image
+            del im
+
+            cells = [
+                OcrCell(
+                    id=ix,
+                    text=line[1],
+                    confidence=line[2],
+                    bbox=BoundingBox.from_tuple(
+                        coord=(
+                            line[0][0][0] / self.scale,
+                            line[0][0][1] / self.scale,
+                            line[0][2][0] / self.scale,
+                            line[0][2][1] / self.scale,
+                        ),
+                        origin=CoordOrigin.TOPLEFT,
+                    ),
+                )
+                for ix, line in enumerate(result)
+            ]
+
+            page.cells = cells  # For now, just overwrites all digital cells.
+
+            # DEBUG code:
+            def draw_clusters_and_cells():
+                image = copy.deepcopy(page.image)
+                draw = ImageDraw.Draw(image)
+
+                cell_color = (
+                    random.randint(30, 140),
+                    random.randint(30, 140),
+                    random.randint(30, 140),
+                )
+                for tc in cells:
+                    x0, y0, x1, y1 = tc.bbox.as_tuple()
+                    draw.rectangle([(x0, y0), (x1, y1)], outline=cell_color)
+                image.show()
+
+            # draw_clusters_and_cells()
+
+            yield page