Fix ocr converter

README.md

@@ -227,6 +227,27 @@ You can also run this via the CLI with
 marker_single FILENAME --use_llm --force_layout_block Table --converter_cls marker.converters.table.TableConverter --output_format json
 ```
 
+### OCR Only
+
+If you only want to run OCR, you can also do that through the `OCRConverter`.
+
+```python
+from marker.converters.ocr import OCRConverter
+from marker.models import create_model_dict
+
+converter = OCRConverter(
+    artifact_dict=create_model_dict(),
+)
+rendered = converter("FILEPATH")
+```
+
+This takes all the same configuration as the PdfConverter.
+
+You can also run this via the CLI with 
+```shell
+marker_single FILENAME --converter_cls marker.converters.ocr.OCRConverter
+```
+
 # Output Formats
 
 ## Markdown

marker/builders/ocr.py

@@ -171,10 +171,12 @@ class OcrBuilder(BaseBuilder):
         before_span, after_span = None, None
         if before_text:
             before_span = copy.deepcopy(span)
+            before_span.structure = []  # Avoid duplicate characters
             before_span.text = before_text
         if after_text:
             after_span = copy.deepcopy(span)
             after_span.text = after_text
+            after_span.structure = []  # Avoid duplicate characters
 
         match_span = copy.deepcopy(span)
         match_span.text = match_text
@@ -214,7 +216,6 @@ class OcrBuilder(BaseBuilder):
                 if not matched:
                     remaining_span = copy.deepcopy(original_span)
                     remaining_span.text = remaining_text
-                    remaining_span.structure = []
                     final_new_spans.append(remaining_span)
                     break
 
@@ -287,10 +288,11 @@ class OcrBuilder(BaseBuilder):
                         current_span.html = (
                             f'<math display="inline">{current_span.text}</math>'
                         )
+
+                    current_chars = self.assign_chars(current_span, current_chars)
                     spans.append(current_span)
                     current_span = None
 
-                    current_chars = self.assign_chars(current_span, current_chars)
                 continue
 
             if not current_span:

marker/converters/ocr.py

@@ -0,0 +1,44 @@
+from typing import Tuple
+
+from marker.builders.document import DocumentBuilder
+from marker.builders.line import LineBuilder
+from marker.builders.ocr import OcrBuilder
+from marker.converters.pdf import PdfConverter
+from marker.processors import BaseProcessor
+from marker.processors.equation import EquationProcessor
+from marker.providers.registry import provider_from_filepath
+from marker.renderers.ocr_json import OCRJSONRenderer
+
+
+class OCRConverter(PdfConverter):
+    default_processors: Tuple[BaseProcessor, ...] = (EquationProcessor,)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if not self.config:
+            self.config = {}
+
+        self.config["format_lines"] = True
+        self.config["keep_chars"] = True
+        self.renderer = OCRJSONRenderer
+
+    def build_document(self, filepath: str):
+        provider_cls = provider_from_filepath(filepath)
+        layout_builder = self.resolve_dependencies(self.layout_builder_class)
+        line_builder = self.resolve_dependencies(LineBuilder)
+        ocr_builder = self.resolve_dependencies(OcrBuilder)
+        document_builder = DocumentBuilder(self.config)
+
+        provider = provider_cls(filepath, self.config)
+        document = document_builder(provider, layout_builder, line_builder, ocr_builder)
+
+        for processor in self.processor_list:
+            processor(document)
+
+        return document
+
+    def __call__(self, filepath: str):
+        document = self.build_document(filepath)
+        renderer = self.resolve_dependencies(self.renderer)
+        return renderer(document)

marker/output.py

@@ -8,6 +8,7 @@ from PIL import Image
 from marker.renderers.html import HTMLOutput
 from marker.renderers.json import JSONOutput, JSONBlockOutput
 from marker.renderers.markdown import MarkdownOutput
+from marker.renderers.ocr_json import OCRJSONOutput
 from marker.schema.blocks import BlockOutput
 from marker.settings import settings
 
@@ -57,6 +58,8 @@ def text_from_rendered(rendered: BaseModel):
         return rendered.html, "html", rendered.images
     elif isinstance(rendered, JSONOutput):
         return rendered.model_dump_json(exclude=["metadata"], indent=2), "json", {}
+    elif isinstance(rendered, OCRJSONOutput):
+        return rendered.model_dump_json(exclude=["metadata"], indent=2), "json", {}
     else:
         raise ValueError("Invalid output type")
 

marker/providers/pdf.py

@@ -239,7 +239,7 @@ class PdfProvider(BaseProvider):
                         )
                         span_chars = [
                             CharClass(
-                                char=c["char"],
+                                text=c["char"],
                                 polygon=PolygonBox.from_bbox(
                                     c["bbox"], ensure_nonzero_area=True
                                 ),

marker/renderers/ocr_json.py

@@ -0,0 +1,128 @@
+from typing import Annotated, List, Tuple
+
+from pydantic import BaseModel
+
+from marker.renderers import BaseRenderer
+from marker.schema import BlockTypes
+from marker.schema.document import Document
+
+
+class OCRJSONCharOutput(BaseModel):
+    id: str
+    block_type: str
+    text: str
+    polygon: List[List[float]]
+    bbox: List[float]
+
+
+class OCRJSONLineOutput(BaseModel):
+    id: str
+    block_type: str
+    html: str
+    polygon: List[List[float]]
+    bbox: List[float]
+    children: List["OCRJSONCharOutput"] | None = None
+
+
+class OCRJSONPageOutput(BaseModel):
+    id: str
+    block_type: str
+    polygon: List[List[float]]
+    bbox: List[float]
+    children: List[OCRJSONLineOutput] | None = None
+
+
+class OCRJSONOutput(BaseModel):
+    children: List[OCRJSONPageOutput]
+    block_type: str = str(BlockTypes.Document)
+    metadata: dict | None = None
+
+
+class OCRJSONRenderer(BaseRenderer):
+    """
+    A renderer for OCR JSON output.
+    """
+
+    image_blocks: Annotated[
+        Tuple[BlockTypes],
+        "The list of block types to consider as images.",
+    ] = (BlockTypes.Picture, BlockTypes.Figure)
+    page_blocks: Annotated[
+        Tuple[BlockTypes],
+        "The list of block types to consider as pages.",
+    ] = (BlockTypes.Page,)
+
+    def extract_json(self, document: Document) -> List[OCRJSONPageOutput]:
+        pages = []
+        for page in document.pages:
+            page_equations = [
+                b for b in page.children if b.block_type == BlockTypes.Equation
+            ]
+            equation_lines = []
+            for equation in page_equations:
+                if not equation.structure:
+                    continue
+
+                equation_lines += [
+                    line
+                    for line in equation.structure
+                    if line.block_type == BlockTypes.Line
+                ]
+
+            page_lines = [
+                block
+                for block in page.children
+                if block.block_type == BlockTypes.Line
+                and block.id not in equation_lines
+            ]
+
+            lines = []
+            for line in page_lines + page_equations:
+                line_obj = OCRJSONLineOutput(
+                    id=str(line.id),
+                    block_type=str(line.block_type),
+                    html="",
+                    polygon=line.polygon.polygon,
+                    bbox=line.polygon.bbox,
+                )
+                if line in page_equations:
+                    line_obj.html = line.html
+                else:
+                    line_obj.html = line.formatted_text(document)
+                    spans = [document.get_block(span_id) for span_id in line.structure]
+                    children = []
+                    for span in spans:
+                        if not span.structure:
+                            continue
+
+                        span_chars = [
+                            document.get_block(char_id) for char_id in span.structure
+                        ]
+                        children.extend(
+                            [
+                                OCRJSONCharOutput(
+                                    id=str(char.id),
+                                    block_type=str(char.block_type),
+                                    text=char.text,
+                                    polygon=char.polygon.polygon,
+                                    bbox=char.polygon.bbox,
+                                )
+                                for char in span_chars
+                            ]
+                        )
+                    line_obj.children = children
+                lines.append(line_obj)
+
+            page = OCRJSONPageOutput(
+                id=str(page.id),
+                block_type=str(page.block_type),
+                polygon=page.polygon.polygon,
+                bbox=page.polygon.bbox,
+                children=lines,
+            )
+            pages.append(page)
+
+        return pages
+
+    def __call__(self, document: Document) -> OCRJSONOutput:
+        return OCRJSONOutput(children=self.extract_json(document), metadata=None)

marker/schema/groups/page.py

@@ -253,14 +253,20 @@ class PageGroup(Group):
                 block.add_structure(line)
                 block.polygon = block.polygon.merge([line.polygon])
                 block.text_extraction_method = text_extraction_method
-                for span in spans:
+                for span_idx, span in enumerate(spans):
                     self.add_full_block(span)
                     line.add_structure(span)
 
                     if not keep_chars:
                         continue
 
-                    for char in provider_output.chars:
+                    # Provider doesn't have chars
+                    if len(provider_output.chars) == 0:
+                        continue
+
+                    # Loop through characters associated with the span
+                    for char in provider_output.chars[span_idx]:
+                        char.page_id = self.page_id
                         self.add_full_block(char)
                         span.add_structure(char)
 

marker/schema/text/char.py

@@ -6,5 +6,5 @@ class Char(Block):
     block_type: BlockTypes = BlockTypes.Char
     block_description: str = "A single character inside a span."
 
-    char: str
+    text: str
     idx: int

tests/conftest.py

@@ -1,4 +1,3 @@
-from marker.providers.pdf import PdfProvider
 import tempfile
 from typing import Dict, Type
 
@@ -19,7 +18,6 @@ from marker.schema.blocks import Block
 from marker.renderers.markdown import MarkdownRenderer
 from marker.renderers.json import JSONRenderer
 from marker.schema.registry import register_block_class
-from marker.services.gemini import GoogleGeminiService
 from marker.util import classes_to_strings, strings_to_classes
 
 
@@ -54,6 +52,7 @@ def table_rec_model(model_dict):
 def ocr_error_model(model_dict):
     yield model_dict["ocr_error_model"]
 
+
 @pytest.fixture(scope="function")
 def config(request):
     config_mark = request.node.get_closest_marker("config")
@@ -65,20 +64,22 @@ def config(request):
 
     return config
 
+
 @pytest.fixture(scope="session")
 def pdf_dataset():
     return datasets.load_dataset("datalab-to/pdfs", split="train")
 
+
 @pytest.fixture(scope="function")
 def temp_doc(request, pdf_dataset):
     filename_mark = request.node.get_closest_marker("filename")
     filename = filename_mark.args[0] if filename_mark else "adversarial.pdf"
 
-    idx = pdf_dataset['filename'].index(filename)
+    idx = pdf_dataset["filename"].index(filename)
     suffix = filename.split(".")[-1]
 
     temp_pdf = tempfile.NamedTemporaryFile(suffix=f".{suffix}")
-    temp_pdf.write(pdf_dataset['pdf'][idx])
+    temp_pdf.write(pdf_dataset["pdf"][idx])
     temp_pdf.flush()
     yield temp_pdf
 
@@ -88,8 +89,17 @@ def doc_provider(request, config, temp_doc):
     provider_cls = provider_from_filepath(temp_doc.name)
     yield provider_cls(temp_doc.name, config)
 
+
 @pytest.fixture(scope="function")
-def pdf_document(request, config, doc_provider, layout_model, ocr_error_model, recognition_model, detection_model):
+def pdf_document(
+    request,
+    config,
+    doc_provider,
+    layout_model,
+    ocr_error_model,
+    recognition_model,
+    detection_model,
+):
     layout_builder = LayoutBuilder(layout_model, config)
     line_builder = LineBuilder(detection_model, ocr_error_model, config)
     ocr_builder = OcrBuilder(recognition_model, config)
@@ -107,7 +117,7 @@ def pdf_converter(request, config, model_dict, renderer, llm_service):
         processor_list=None,
         renderer=classes_to_strings([renderer])[0],
         config=config,
-        llm_service=llm_service
+        llm_service=llm_service,
     )
 
 

tests/converters/test_ocr_converter.py

@@ -0,0 +1,19 @@
+import pytest
+
+from marker.converters.ocr import OCRConverter
+from marker.renderers.ocr_json import OCRJSONOutput
+
+
+def _ocr_converter(config, model_dict, temp_pdf):
+    converter = OCRConverter(artifact_dict=model_dict, config=config)
+
+    ocr_json: OCRJSONOutput = converter(temp_pdf.name)
+    pages = ocr_json.pages
+
+    assert len(pages) == 1
+    breakpoint()
+
+
+@pytest.mark.config({"page_range": [0]})
+def test_ocr_converter(config, model_dict, temp_doc):
+    _ocr_converter(config, model_dict, temp_doc)