Merge remote-tracking branch 'origin/vik_v2' into dev-mose/marker-v2

marker/v2/builders/structure.py

@@ -12,7 +12,7 @@ from marker.v2.schema.groups.page import PageGroup
 class StructureBuilder(BaseBuilder):
     gap_threshold: int = 10
 
-    def __init__(self, config: Optional[BaseModel] = None):
+    def __init__(self, config=None):
         super().__init__(config)
 
     def __call__(self, document: Document):

marker/v2/converters/__init__.py

@@ -8,6 +8,7 @@ from marker.v2.util import assign_config
 class BaseConverter:
     def __init__(self, config: Optional[BaseModel | dict] = None):
         assign_config(self, config)
+        self.config = config
 
     def __call__(self, *args, **kwargs):
         raise NotImplementedError

marker/v2/converters/pdf.py

@@ -1,9 +1,8 @@
-<<<<<<< HEAD
 import os
-=======
-from marker.v2.providers.pdf import PdfProvider
+os.environ["TOKENIZERS_PARALLELISM"] = "false" # disables a tokenizers warning
 
->>>>>>> origin/v2
+from marker.v2.processors.sectionheader import SectionHeaderProcessor
+from marker.v2.providers.pdf import PdfProvider
 import tempfile
 from typing import List, Optional
 
@@ -24,7 +23,7 @@ from marker.v2.renderers.markdown import MarkdownRenderer
 
 
 class PdfConverter(BaseConverter):
-    def __init__(self, config: Optional[BaseModel] = None):
+    def __init__(self, config=None):
         super().__init__(config)
 
         self.layout_model = setup_layout_model()
@@ -33,21 +32,24 @@ class PdfConverter(BaseConverter):
         self.table_rec_model = setup_table_rec_model()
         self.detection_model = setup_detection_model()
 
-    def __call__(self, filepath: str, page_range: List[int] | None = None):
-        pdf_provider = PdfProvider(filepath, {"page_range": page_range, "force_ocr": False})
+    def __call__(self, filepath: str):
+        pdf_provider = PdfProvider(filepath, self.config)
 
-        layout_builder = LayoutBuilder(self.layout_model)
-        ocr_builder = OcrBuilder(self.detection_model, self.recognition_model)
-        document = DocumentBuilder()(pdf_provider, layout_builder, ocr_builder)
-        StructureBuilder()(document)
+        layout_builder = LayoutBuilder(self.layout_model, self.config)
+        ocr_builder = OcrBuilder(self.detection_model, self.recognition_model, self.config)
+        document = DocumentBuilder(self.config)(pdf_provider, layout_builder, ocr_builder)
+        StructureBuilder(self.config)(document)
 
-        equation_processor = EquationProcessor(self.texify_model)
+        equation_processor = EquationProcessor(self.texify_model, self.config)
         equation_processor(document)
 
-        table_processor = TableProcessor(self.detection_model, self.recognition_model, self.table_rec_model)
+        table_processor = TableProcessor(self.detection_model, self.recognition_model, self.table_rec_model, self.config)
         table_processor(document)
 
-        renderer = MarkdownRenderer()
+        section_header_processor = SectionHeaderProcessor(self.config)
+        section_header_processor(document)
+
+        renderer = MarkdownRenderer(self.config)
         return renderer(document)
 
 

marker/v2/processors/equation.py

@@ -16,7 +16,7 @@ class EquationProcessor(BaseProcessor):
     batch_size = None
     token_buffer = 256
 
-    def __init__(self, texify_model, config: Optional[BaseModel] = None):
+    def __init__(self, texify_model, config=None):
         super().__init__(config)
 
         self.texify_model = texify_model

marker/v2/processors/sectionheader.py

@@ -0,0 +1,84 @@
+from marker.v2.processors import BaseProcessor
+from marker.v2.schema import BlockTypes
+from marker.v2.schema.document import Document
+
+import numpy as np
+from sklearn.cluster import KMeans
+from sklearn.exceptions import ConvergenceWarning
+
+# Ignore sklearn warning about not converging
+import warnings
+warnings.filterwarnings("ignore", category=ConvergenceWarning)
+
+
+class SectionHeaderProcessor(BaseProcessor):
+    block_types = (BlockTypes.SectionHeader, )
+    level_count = 4
+    merge_threshold = .25
+    default_level = 2
+    height_tolerance = .99
+
+    def __call__(self, document: Document):
+        line_heights = {}
+        for page in document.pages:
+            for block in page.children:
+                if block.block_type not in self.block_types:
+                    continue
+
+                line_heights[block.block_id] = [document.get_block(l).polygon.height for l in block.structure if l.block_type == BlockTypes.Line]
+
+        flat_line_heights = [h for heights in line_heights.values() for h in heights]
+        heading_ranges = self.bucket_headings(flat_line_heights)
+
+        for page in document.pages:
+            for block in page.children:
+                if block.block_type not in self.block_types:
+                    continue
+
+                block_heights = line_heights[block.block_id]
+                if len(block_heights) > 0:
+                    avg_height = sum(block_heights) / len(block_heights)
+                    for idx, (min_height, max_height) in enumerate(heading_ranges):
+                        if avg_height >= min_height * self.height_tolerance:
+                            block.heading_level = idx + 1
+                            break
+
+                if block.heading_level is None:
+                    block.heading_level = self.default_level
+
+    def bucket_headings(self, line_heights, num_levels=4):
+        if len(line_heights) <= self.level_count:
+            return []
+
+        data = np.asarray(line_heights).reshape(-1, 1)
+        labels = KMeans(n_clusters=num_levels, random_state=0, n_init="auto").fit_predict(data)
+        data_labels = np.concatenate([data, labels.reshape(-1, 1)], axis=1)
+        data_labels = np.sort(data_labels, axis=0)
+
+        cluster_means = {int(label): float(np.mean(data_labels[data_labels[:, 1] == label, 0])) for label in np.unique(labels)}
+        label_max = None
+        label_min = None
+        heading_ranges = []
+        prev_cluster = None
+        for row in data_labels:
+            value, label = row
+            value = float(value)
+            label = int(label)
+            if prev_cluster is not None and label != prev_cluster:
+                prev_cluster_mean = cluster_means[prev_cluster]
+                cluster_mean = cluster_means[label]
+                if cluster_mean * self.merge_threshold < prev_cluster_mean:
+                    heading_ranges.append((label_min, label_max))
+                    label_min = None
+                    label_max = None
+
+            label_min = value if label_min is None else min(label_min, value)
+            label_max = value if label_max is None else max(label_max, value)
+            prev_cluster = label
+
+        if label_min is not None:
+            heading_ranges.append((label_min, label_max))
+
+        heading_ranges = sorted(heading_ranges, reverse=True)
+
+        return heading_ranges

marker/v2/processors/table.py

@@ -18,7 +18,7 @@ class TableProcessor(BaseProcessor):
     table_rec_batch_size = None
     ocr_batch_size = None
 
-    def __init__(self, detection_model, ocr_model, table_rec_model, config: Optional[BaseModel] = None):
+    def __init__(self, detection_model, ocr_model, table_rec_model, config=None):
         super().__init__(config)
 
         self.detection_model = detection_model

marker/v2/providers/pdf.py

@@ -23,7 +23,7 @@ class PdfProvider(BaseProvider):
     flatten_pdf: bool = True
     force_ocr: bool = False
 
-    def __init__(self, filepath: str, config: Optional[BaseModel] = None):
+    def __init__(self, filepath: str, config = None):
         super().__init__(filepath, config)
 
         self.doc: pdfium.PdfDocument = pdfium.PdfDocument(self.filepath)

marker/v2/renderers/__init__.py

@@ -3,16 +3,14 @@ from typing import Optional
 from pydantic import BaseModel
 
 from marker.v2.schema import BlockTypes
-
+from marker.v2.util import assign_config
 
 
 class BaseRenderer:
     block_type: BlockTypes | None = None
 
     def __init__(self, config: Optional[BaseModel | dict] = None):
-        if config:
-            for k in config.model_fields:
-                setattr(self, k, config[k])
+        assign_config(self, config)
 
     def __call__(self, document):
         # Children are in reading order

marker/v2/renderers/html.py

@@ -1,12 +1,16 @@
 import re
 
-from bs4 import BeautifulSoup
+from bs4 import BeautifulSoup, MarkupResemblesLocatorWarning
 from pydantic import BaseModel
 
 from marker.v2.renderers import BaseRenderer
 from marker.v2.schema import BlockTypes
 from marker.v2.schema.blocks import BlockId
 
+# Ignore beautifulsoup warnings
+import warnings
+warnings.filterwarnings("ignore", category=MarkupResemblesLocatorWarning)
+
 
 class HTMLOutput(BaseModel):
     html: str
@@ -34,6 +38,8 @@ def merge_consecutive_tags(html, tag):
 class HTMLRenderer(BaseRenderer):
     remove_blocks: list = [BlockTypes.PageHeader, BlockTypes.PageFooter]
     image_blocks: list = [BlockTypes.Picture, BlockTypes.Figure]
+    page_blocks: list = [BlockTypes.Page]
+    paginate_output: bool = False
 
     def extract_image(self, document, image_id):
         image_block = document.get_block(image_id)
@@ -65,6 +71,11 @@ class HTMLRenderer(BaseRenderer):
                 image_name = f"{ref_block_id.to_path()}.png"
                 images[image_name] = image
                 ref.replace_with(BeautifulSoup(f"<p><img src='{image_name}'></p>", 'html.parser'))
+            elif ref_block_id.block_type in self.page_blocks:
+                images.update(sub_images)
+                if self.paginate_output:
+                    content = f"<div class='page' data-page-id='{ref_block_id.page_id}'>{content}</div>"
+                ref.replace_with(BeautifulSoup(f"{content}", 'html.parser'))
             else:
                 images.update(sub_images)
                 ref.replace_with(BeautifulSoup(f"{content}", 'html.parser'))

marker/v2/renderers/markdown.py

@@ -6,7 +6,19 @@ from marker.v2.schema.document import Document
 
 
 class Markdownify(MarkdownConverter):
-    pass
+    def __init__(self, paginate_output, page_separator, **kwargs):
+        super().__init__(**kwargs)
+        self.paginate_output = paginate_output
+        self.page_separator = page_separator
+
+    def convert_div(self, el, text, convert_as_inline):
+        is_page = el.has_attr('class') and el['class'][0] == 'page'
+        if self.paginate_output and is_page:
+            page_id = el['data-page-id']
+            pagination_item = "\n\n" + "{" + str(page_id) + "}" + self.page_separator + "\n\n"
+            return pagination_item + text
+        else:
+            return text
 
 
 class MarkdownOutput(BaseModel):
@@ -15,10 +27,14 @@ class MarkdownOutput(BaseModel):
 
 
 class MarkdownRenderer(HTMLRenderer):
+    page_separator: str = "-" * 48
+
     def __call__(self, document: Document) -> MarkdownOutput:
         document_output = document.render()
         full_html, images = self.extract_html(document, document_output)
         md_cls = Markdownify(
+            self.paginate_output,
+            self.page_separator,
             heading_style="ATX",
             bullets="-",
             escape_misc=False,

marker/v2/schema/blocks/base.py

@@ -28,6 +28,9 @@ class BlockId(BaseModel):
             return f"/page/{self.page_id}"
         return f"/page/{self.page_id}/{self.block_type.name}/{self.block_id}"
 
+    def __hash__(self):
+        return hash(str(self))
+
     def __repr__(self):
         return str(self)
 

marker/v2/schema/blocks/sectionheader.py

@@ -4,8 +4,10 @@ from marker.v2.schema.blocks import Block
 
 class SectionHeader(Block):
     block_type: BlockTypes = BlockTypes.SectionHeader
+    heading_level: int | None = None
 
     def assemble_html(self, child_blocks, parent_structure):
         template = super().assemble_html(child_blocks, parent_structure)
         template = template.replace("\n", " ")
-        return f"<h2>{template}</h2>"
+        tag = f"h{self.heading_level}" if self.heading_level else "h2"
+        return f"<{tag}>{template}</{tag}>"

marker/v2/schema/groups/page.py

@@ -48,6 +48,12 @@ class PageGroup(Block):
         assert block.block_id == block_id.block_id
         return block
 
+    def assemble_html(self, child_blocks, parent_structure=None):
+        template = ""
+        for c in child_blocks:
+            template += f"<content-ref src='{c.id}'></content-ref>"
+        return template
+
     def merge_blocks(
         self,
         page_lines: List[Line],

marker/v2/util.py

@@ -2,11 +2,24 @@ from pydantic import BaseModel
 
 
 def assign_config(cls, config: BaseModel | dict | None):
+    cls_name = cls.__class__.__name__
     if config is None:
         return
     elif isinstance(config, BaseModel):
-        for k in config.model_fields:
-            setattr(cls, k, config[k])
+        dict_config = config.dict()
     elif isinstance(config, dict):
-        for k, v in config.items():
-            setattr(cls, k, v)
+        dict_config = config
+    else:
+        raise ValueError("config must be a dict or a pydantic BaseModel")
+
+    for k in dict_config:
+        if hasattr(cls, k):
+            setattr(cls, k, dict_config[k])
+    for k in dict_config:
+        if cls_name not in k:
+            continue
+        # Enables using class-specific keys, like "MarkdownRenderer_remove_blocks"
+        split_k = k.removeprefix(cls_name + "_")
+
+        if hasattr(cls, split_k):
+            setattr(cls, split_k, dict_config[k])

tests/utils.py

@@ -11,10 +11,7 @@ from marker.v2.schema.document import Document
 
 def setup_pdf_document(
     filename='adversarial.pdf',
-    pdf_provider_config=None,
-    layout_builder_config=None,
-    ocr_builder_config=None,
-    document_builder_config=None
+    config=None,
 ) -> Document:
     dataset = datasets.load_dataset("datalab-to/pdfs", split="train")
     idx = dataset['filename'].index(filename)
@@ -27,9 +24,9 @@ def setup_pdf_document(
     recognition_model = setup_recognition_model()
     detection_model = setup_detection_model()
 
-    provider = PdfProvider(temp_pdf.name, pdf_provider_config)
-    layout_builder = LayoutBuilder(layout_model, layout_builder_config)
-    ocr_builder = OcrBuilder(detection_model, recognition_model, ocr_builder_config)
-    builder = DocumentBuilder(document_builder_config)
+    provider = PdfProvider(temp_pdf.name, config)
+    layout_builder = LayoutBuilder(layout_model, config)
+    ocr_builder = OcrBuilder(detection_model, recognition_model, config)
+    builder = DocumentBuilder(config)
     document = builder(provider, layout_builder, ocr_builder)
     return document