Merge remote-tracking branch 'origin/dev' into highquality-processors

marker/builders/layout.py

@@ -5,6 +5,10 @@ from surya.layout import batch_layout_detection
 from surya.schema import LayoutResult
 from surya.model.layout.encoderdecoder import SuryaLayoutModel
 
+from surya.ocr_error import batch_ocr_error_detection
+from surya.schema import OCRErrorDetectionResult
+from surya.model.ocr_error.model import DistilBertForSequenceClassification
+
 from marker.settings import settings
 from marker.builders import BaseBuilder
 from marker.providers import ProviderOutput, ProviderPageLines
@@ -37,15 +41,21 @@ class LayoutBuilder(BaseBuilder):
         document_ocr_threshold (float):
             The minimum ratio of pages that must pass the layout coverage check
             to avoid OCR. Default is 0.8.
+
+        error_model_segment_length (int):
+            The maximum number of characters to send to the OCR error model.
+            Default is 1024.
     """
     batch_size = None
     layout_coverage_min_lines = 1
     layout_coverage_threshold = .1
     document_ocr_threshold = .8
+    error_model_segment_length = 512
     excluded_for_coverage = (BlockTypes.Figure, BlockTypes.Picture, BlockTypes.Table, BlockTypes.FigureGroup, BlockTypes.TableGroup, BlockTypes.PictureGroup)
 
-    def __init__(self, layout_model: SuryaLayoutModel, config=None):
+    def __init__(self, layout_model: SuryaLayoutModel, ocr_error_model: DistilBertForSequenceClassification, config=None):
         self.layout_model = layout_model
+        self.ocr_error_model = ocr_error_model
 
         super().__init__(config)
 
@@ -71,6 +81,31 @@ class LayoutBuilder(BaseBuilder):
         )
         return layout_results
 
+    def surya_ocr_error_detection(self, pages:List[PageGroup], provider_page_lines: ProviderPageLines) -> OCRErrorDetectionResult:
+        page_texts = []
+        for document_page in pages:
+            page_text = ''
+            provider_lines = provider_page_lines.get(document_page.page_id, [])
+            for line in provider_lines:
+                page_text += ' '.join([s.text for s in line.spans])
+
+            # Sample text from the middle
+            if len(page_text) > 0:
+                page_text_middle = len(page_text) // 2
+                page_text_start = max(0, page_text_middle - self.error_model_segment_length // 2)
+                page_text_end = page_text_start + self.error_model_segment_length
+                page_text = page_text[page_text_start:page_text_end]
+
+            page_texts.append(page_text)
+
+        ocr_error_detection_results = batch_ocr_error_detection(
+            page_texts,
+            self.ocr_error_model,
+            self.ocr_error_model.tokenizer,
+            batch_size=int(self.get_batch_size())       #TODO Better Multiplier
+        )
+        return ocr_error_detection_results
+
     def add_blocks_to_pages(self, pages: List[PageGroup], layout_results: List[LayoutResult]):
         for page, layout_result in zip(pages, layout_results):
             layout_page_size = PolygonBox.from_bbox(layout_result.image_bbox).size
@@ -92,16 +127,17 @@ class LayoutBuilder(BaseBuilder):
                 page.children = []
 
     def merge_blocks(self, document_pages: List[PageGroup], provider_page_lines: ProviderPageLines):
+        ocr_error_detection_labels = self.surya_ocr_error_detection(document_pages, provider_page_lines).labels
+
         good_pages = []
-        for document_page in document_pages:
+        for (document_page, ocr_error_detection_label) in zip(document_pages, ocr_error_detection_labels):
             provider_lines = provider_page_lines.get(document_page.page_id, [])
-            good_pages.append(self.check_layout_coverage(document_page, provider_lines))
+            good_pages.append(self.check_layout_coverage(document_page, provider_lines) and (ocr_error_detection_label != "bad"))
 
         ocr_document = sum(good_pages) / len(good_pages) < self.document_ocr_threshold
         for idx, document_page in enumerate(document_pages):
             provider_lines = provider_page_lines.get(document_page.page_id, [])
             needs_ocr = not good_pages[idx]
-
             if needs_ocr and ocr_document:
                 document_page.text_extraction_method = "surya"
                 continue

marker/models.py

@@ -12,12 +12,15 @@ from surya.model.recognition.model import load_model as load_recognition_model
 from surya.model.recognition.processor import load_processor as load_recognition_processor
 from surya.model.table_rec.model import load_model as load_table_model
 from surya.model.table_rec.processor import load_processor as load_table_processor
+from surya.model.ocr_error.model import load_model as load_ocr_error_model
+from surya.model.ocr_error.model import load_tokenizer as load_ocr_error_tokenizer
 
 from texify.model.model import GenerateVisionEncoderDecoderModel
 from surya.model.layout.encoderdecoder import SuryaLayoutModel
 from surya.model.detection.model import EfficientViTForSemanticSegmentation
 from surya.model.recognition.encoderdecoder import OCREncoderDecoderModel
 from surya.model.table_rec.encoderdecoder import TableRecEncoderDecoderModel
+from surya.model.ocr_error.model import DistilBertForSequenceClassification
 
 
 def setup_table_rec_model(device=None, dtype=None) -> TableRecEncoderDecoderModel:
@@ -64,6 +67,13 @@ def setup_layout_model(device=None, dtype=None) -> SuryaLayoutModel:
     model.processor = load_layout_processor()
     return model
 
+def setup_ocr_error_model(device=None, dtype=None) -> DistilBertForSequenceClassification:
+    if device:
+        model = load_ocr_error_model(device=device, dtype=dtype)
+    else:
+        model = load_ocr_error_model()
+    model.tokenizer = load_ocr_error_tokenizer()
+    return model
 
 def create_model_dict(device=None, dtype=None) -> dict:
     return {
@@ -72,4 +82,5 @@ def create_model_dict(device=None, dtype=None) -> dict:
         "recognition_model": setup_recognition_model(device, dtype),
         "table_rec_model": setup_table_rec_model(device, dtype),
         "detection_model": setup_detection_model(device, dtype),
+        "ocr_error_model": setup_ocr_error_model(device,dtype)
     }

marker/renderers/markdown.py

@@ -53,6 +53,15 @@ class Markdownify(MarkdownConverter):
         else:
             return "\n" + self.block_math_delimiters[0] + text + self.block_math_delimiters[1] + "\n\n"
 
+    def convert_td(self, el, text, convert_as_inline):
+        text = text.replace("|", " ").replace("\n", " ")
+        return super().convert_td(el, text, convert_as_inline)
+
+    def convert_th(self, el, text, convert_as_inline):
+        text = text.replace("|", " ").replace("\n", " ")
+        return super().convert_th(el, text, convert_as_inline)
+
+
 
 class MarkdownOutput(BaseModel):
     markdown: str

poetry.lock

@@ -1447,13 +1447,13 @@ test = ["flaky", "ipyparallel", "pre-commit", "pytest (>=7.0)", "pytest-asyncio
 
 [[package]]
 name = "ipython"
-version = "8.30.0"
+version = "8.31.0"
 description = "IPython: Productive Interactive Computing"
 optional = false
 python-versions = ">=3.10"
 files = [
-    {file = "ipython-8.30.0-py3-none-any.whl", hash = "sha256:85ec56a7e20f6c38fce7727dcca699ae4ffc85985aa7b23635a8008f918ae321"},
-    {file = "ipython-8.30.0.tar.gz", hash = "sha256:cb0a405a306d2995a5cbb9901894d240784a9f341394c6ba3f4fe8c6eb89ff6e"},
+    {file = "ipython-8.31.0-py3-none-any.whl", hash = "sha256:46ec58f8d3d076a61d128fe517a51eb730e3aaf0c184ea8c17d16e366660c6a6"},
+    {file = "ipython-8.31.0.tar.gz", hash = "sha256:b6a2274606bec6166405ff05e54932ed6e5cfecaca1fc05f2cacde7bb074d70b"},
 ]
 
 [package.dependencies]
@@ -1759,13 +1759,13 @@ jupyter-server = ">=1.1.2"
 
 [[package]]
 name = "jupyter-server"
-version = "2.14.2"
+version = "2.15.0"
 description = "The backend—i.e. core services, APIs, and REST endpoints—to Jupyter web applications."
 optional = false
-python-versions = ">=3.8"
+python-versions = ">=3.9"
 files = [
-    {file = "jupyter_server-2.14.2-py3-none-any.whl", hash = "sha256:47ff506127c2f7851a17bf4713434208fc490955d0e8632e95014a9a9afbeefd"},
-    {file = "jupyter_server-2.14.2.tar.gz", hash = "sha256:66095021aa9638ced276c248b1d81862e4c50f292d575920bbe960de1c56b12b"},
+    {file = "jupyter_server-2.15.0-py3-none-any.whl", hash = "sha256:872d989becf83517012ee669f09604aa4a28097c0bd90b2f424310156c2cdae3"},
+    {file = "jupyter_server-2.15.0.tar.gz", hash = "sha256:9d446b8697b4f7337a1b7cdcac40778babdd93ba614b6d68ab1c0c918f1c4084"},
 ]
 
 [package.dependencies]
@@ -1774,7 +1774,7 @@ argon2-cffi = ">=21.1"
 jinja2 = ">=3.0.3"
 jupyter-client = ">=7.4.4"
 jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0"
-jupyter-events = ">=0.9.0"
+jupyter-events = ">=0.11.0"
 jupyter-server-terminals = ">=0.4.4"
 nbconvert = ">=6.4.4"
 nbformat = ">=5.3.0"
@@ -3392,24 +3392,24 @@ diagrams = ["jinja2", "railroad-diagrams"]
 
 [[package]]
 name = "pypdfium2"
-version = "4.30.0"
+version = "4.30.1"
 description = "Python bindings to PDFium"
 optional = false
 python-versions = ">=3.6"
 files = [
-    {file = "pypdfium2-4.30.0-py3-none-macosx_10_13_x86_64.whl", hash = "sha256:b33ceded0b6ff5b2b93bc1fe0ad4b71aa6b7e7bd5875f1ca0cdfb6ba6ac01aab"},
-    {file = "pypdfium2-4.30.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:4e55689f4b06e2d2406203e771f78789bd4f190731b5d57383d05cf611d829de"},
-    {file = "pypdfium2-4.30.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e6e50f5ce7f65a40a33d7c9edc39f23140c57e37144c2d6d9e9262a2a854854"},
-    {file = "pypdfium2-4.30.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3d0dd3ecaffd0b6dbda3da663220e705cb563918249bda26058c6036752ba3a2"},
-    {file = "pypdfium2-4.30.0-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cc3bf29b0db8c76cdfaac1ec1cde8edf211a7de7390fbf8934ad2aa9b4d6dfad"},
-    {file = "pypdfium2-4.30.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f1f78d2189e0ddf9ac2b7a9b9bd4f0c66f54d1389ff6c17e9fd9dc034d06eb3f"},
-    {file = "pypdfium2-4.30.0-py3-none-musllinux_1_1_aarch64.whl", hash = "sha256:5eda3641a2da7a7a0b2f4dbd71d706401a656fea521b6b6faa0675b15d31a163"},
-    {file = "pypdfium2-4.30.0-py3-none-musllinux_1_1_i686.whl", hash = "sha256:0dfa61421b5eb68e1188b0b2231e7ba35735aef2d867d86e48ee6cab6975195e"},
-    {file = "pypdfium2-4.30.0-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:f33bd79e7a09d5f7acca3b0b69ff6c8a488869a7fab48fdf400fec6e20b9c8be"},
-    {file = "pypdfium2-4.30.0-py3-none-win32.whl", hash = "sha256:ee2410f15d576d976c2ab2558c93d392a25fb9f6635e8dd0a8a3a5241b275e0e"},
-    {file = "pypdfium2-4.30.0-py3-none-win_amd64.whl", hash = "sha256:90dbb2ac07be53219f56be09961eb95cf2473f834d01a42d901d13ccfad64b4c"},
-    {file = "pypdfium2-4.30.0-py3-none-win_arm64.whl", hash = "sha256:119b2969a6d6b1e8d55e99caaf05290294f2d0fe49c12a3f17102d01c441bd29"},
-    {file = "pypdfium2-4.30.0.tar.gz", hash = "sha256:48b5b7e5566665bc1015b9d69c1ebabe21f6aee468b509531c3c8318eeee2e16"},
+    {file = "pypdfium2-4.30.1-py3-none-macosx_10_13_x86_64.whl", hash = "sha256:e07c47633732cc18d890bb7e965ad28a9c5a932e548acb928596f86be2e5ae37"},
+    {file = "pypdfium2-4.30.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:5ea2d44e96d361123b67b00f527017aa9c847c871b5714e013c01c3eb36a79fe"},
+    {file = "pypdfium2-4.30.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1de7a3a36803171b3f66911131046d65a732f9e7834438191cb58235e6163c4e"},
+    {file = "pypdfium2-4.30.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b8a4231efb13170354f568c722d6540b8d5b476b08825586d48ef70c40d16e03"},
+    {file = "pypdfium2-4.30.1-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6f434a4934e8244aa95343ffcf24e9ad9f120dbb4785f631bb40a88c39292493"},
+    {file = "pypdfium2-4.30.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f454032a0bc7681900170f67d8711b3942824531e765f91c2f5ce7937f999794"},
+    {file = "pypdfium2-4.30.1-py3-none-musllinux_1_1_aarch64.whl", hash = "sha256:bbf9130a72370ee9d602e39949b902db669a2a1c24746a91e5586eb829055d9f"},
+    {file = "pypdfium2-4.30.1-py3-none-musllinux_1_1_i686.whl", hash = "sha256:5cb52884b1583b96e94fd78542c63bb42e06df5e8f9e52f8f31f5ad5a1e53367"},
+    {file = "pypdfium2-4.30.1-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:1a9e372bd4867ff223cc8c338e33fe11055dad12f22885950fc27646cc8d9122"},
+    {file = "pypdfium2-4.30.1-py3-none-win32.whl", hash = "sha256:421f1cf205e213e07c1f2934905779547f4f4a2ff2f59dde29da3d511d3fc806"},
+    {file = "pypdfium2-4.30.1-py3-none-win_amd64.whl", hash = "sha256:598a7f20264ab5113853cba6d86c4566e4356cad037d7d1f849c8c9021007e05"},
+    {file = "pypdfium2-4.30.1-py3-none-win_arm64.whl", hash = "sha256:c2b6d63f6d425d9416c08d2511822b54b8e3ac38e639fc41164b1d75584b3a8c"},
+    {file = "pypdfium2-4.30.1.tar.gz", hash = "sha256:5f5c7c6d03598e107d974f66b220a49436aceb191da34cda5f692be098a814ce"},
 ]
 
 [[package]]

tests/builders/test_blank_page.py

@@ -5,8 +5,8 @@ from marker.builders.layout import LayoutBuilder
 from marker.builders.ocr import OcrBuilder
 
 
-def test_blank_page(config, pdf_provider, layout_model, recognition_model, detection_model):
-    layout_builder = LayoutBuilder(layout_model, config)
+def test_blank_page(config, pdf_provider, layout_model, ocr_error_model, recognition_model, detection_model):
+    layout_builder = LayoutBuilder(layout_model, ocr_error_model, config)
     builder = DocumentBuilder(config)
     document = builder.build_document(pdf_provider)
 

tests/builders/test_garbled_pdf.py

@@ -1,8 +1,9 @@
 import pytest
-from marker.schema import BlockTypes
 
+from marker.builders.document import DocumentBuilder
+from marker.builders.layout import LayoutBuilder
+from marker.schema import BlockTypes
 
-@pytest.mark.skip(reason="This is failing because we need better garbled text detection")
 @pytest.mark.filename("water_damage.pdf")
 def test_garbled_pdf(pdf_document):
     assert pdf_document.pages[0].structure[0] == '/page/0/Table/0'
@@ -18,3 +19,29 @@ def test_garbled_pdf(pdf_document):
     span = pdf_document.pages[0].get_block(table_cell.structure[0])
     assert span.block_type == BlockTypes.Span
     assert "комплекс" in span.text
+
+
+@pytest.mark.filename("hindi_judgement.pdf")
+@pytest.mark.config({"page_range": [2, 3]})
+def test_garbled_builder(config, pdf_provider, layout_model, ocr_error_model):
+    layout_builder = LayoutBuilder(layout_model, ocr_error_model, config)
+    builder = DocumentBuilder(config)
+    document = builder.build_document(pdf_provider)
+
+    bad_ocr_results = layout_builder.surya_ocr_error_detection(document.pages, pdf_provider.page_lines)
+    assert len(bad_ocr_results.labels) == 2
+    assert all([l == "bad" for l in bad_ocr_results.labels])
+
+
+@pytest.mark.filename("adversarial.pdf")
+@pytest.mark.config({"page_range": [2, 3]})
+def test_nongarbled_builder(config, pdf_provider, layout_model, ocr_error_model):
+    layout_builder = LayoutBuilder(layout_model, ocr_error_model, config)
+    builder = DocumentBuilder(config)
+    document = builder.build_document(pdf_provider)
+
+    bad_ocr_results = layout_builder.surya_ocr_error_detection(document.pages, pdf_provider.page_lines)
+    assert len(bad_ocr_results.labels) == 2
+    assert all([l == "good" for l in bad_ocr_results.labels])
+
+

tests/conftest.py

@@ -11,7 +11,7 @@ from marker.builders.ocr import OcrBuilder
 from marker.converters.pdf import PdfConverter
 from marker.models import setup_detection_model, setup_layout_model, \
     setup_recognition_model, setup_table_rec_model, \
-    setup_texify_model
+    setup_texify_model, setup_ocr_error_model
 from marker.schema import BlockTypes
 from marker.schema.blocks import Block
 from marker.renderers.markdown import MarkdownRenderer
@@ -54,6 +54,11 @@ def table_rec_model():
     yield table_rec_m
     del table_rec_m
 
+@pytest.fixture(scope="session")
+def ocr_error_model():
+    ocr_error_m = setup_ocr_error_model()
+    yield ocr_error_m
+    del ocr_error_m
 
 @pytest.fixture(scope="function")
 def config(request):
@@ -87,8 +92,8 @@ def pdf_provider(request, config, temp_pdf):
 
 
 @pytest.fixture(scope="function")
-def pdf_document(request, config, pdf_provider, layout_model, recognition_model, detection_model):
-    layout_builder = LayoutBuilder(layout_model, config)
+def pdf_document(request, config, pdf_provider, layout_model, ocr_error_model, recognition_model, detection_model):
+    layout_builder = LayoutBuilder(layout_model, ocr_error_model, config)
     ocr_builder = OcrBuilder(detection_model, recognition_model, config)
     builder = DocumentBuilder(config)
     document = builder(pdf_provider, layout_builder, ocr_builder)
@@ -96,13 +101,14 @@ def pdf_document(request, config, pdf_provider, layout_model, recognition_model,
 
 
 @pytest.fixture(scope="function")
-def pdf_converter(request, config, layout_model, texify_model, recognition_model, table_rec_model, detection_model, renderer):
+def pdf_converter(request, config, layout_model, texify_model, recognition_model, table_rec_model, detection_model, ocr_error_model, renderer):
     model_dict = {
         "layout_model": layout_model,
         "texify_model": texify_model,
         "recognition_model": recognition_model,
         "table_rec_model": table_rec_model,
-        "detection_model": detection_model
+        "detection_model": detection_model,
+        "ocr_error_model": ocr_error_model
     }
     yield PdfConverter(
         artifact_dict=model_dict,