Add structure

marker/v2/builders/__init__.py

@@ -1,6 +1,8 @@
 class BaseBuilder:
     def __init__(self, config=None):
-        self.config = config
+        if config:
+            for k in config:
+                setattr(self, k, config[k])
 
     def __call__(self, data, *args, **kwargs):
         raise NotImplementedError

marker/v2/builders/structure.py

@@ -1,53 +1,85 @@
+from typing import List, Optional
+
 from pydantic import BaseModel
 
 from marker.v2.builders import BaseBuilder
 from marker.v2.schema.document import Document
-from marker.v2.schema.groups import GROUP_BLOCK_REGISTRY
+from marker.v2.schema.groups import GROUP_BLOCK_REGISTRY, ListGroup
+from marker.v2.schema.groups.page import PageGroup
 
 
-class StructureConfig(BaseModel):
+class StructureBuilder(BaseBuilder):
     gap_threshold: int = 10
 
-
-class StructureBuilder(BaseBuilder):
-    def __init__(self, config):
+    def __init__(self, config: Optional[BaseModel] = None):
         super().__init__(config)
 
     def __call__(self, document: Document):
         for page in document.pages:
-            initial_structure = [block._id for block in page.children]
+            page.structure = [block._id for block in page.children]
+            self.group_caption_blocks(page)
+            self.group_lists(page)
 
-    def group_caption_blocks(self, page):
+    def group_caption_blocks(self, page: PageGroup):
         for i, block in enumerate(page.children):
-            if block.block_type in ["Table", "Figure", "Picture"]:
-                block_structure = [block._id]
-
-                for j, prev_block in enumerate(page.children[:i][::-1]):
-                    if all([
-                        prev_block.block_type in ["Caption", "Footnote"],
-                        prev_block.minimum_gap(block) < self.config.gap_threshold
-                    ]):
-                        block_structure.append(prev_block._id)
-                    else:
-                        break
-
-                for j, next_block in enumerate(page.children[i + 1:]):
-                    if all([
-                        next_block.block_type in ["Caption", "Footnote"],
-                        next_block.minimum_gap(block) < self.config.gap_threshold
-                    ]):
-                        block_structure.append(next_block._id)
-                    else:
-                        break
-
-                if len(block_structure) > 0:
-                    new_block_cls = GROUP_BLOCK_REGISTRY[block.block_type]
-                    # TODO: fix the polygon to span all the blocks inside the grouped block
-                    # TODO: Add the structure, etc, to the block
-                    block = page.add_block(new_block_cls, block.polygon)
-
-        # Table, Figure, Picture
-        pass
-
-    def group_lists(self, page):
-        pass
+            if block.block_type not in ["Table", "Figure", "Picture"]:
+                continue
+
+            block_structure = [block._id]
+            selected_polygons = [block.polygon]
+
+            for j, prev_block in enumerate(page.children[:i][::-1]):
+                if all([
+                    prev_block.block_type in ["Caption", "Footnote"],
+                    prev_block.minimum_gap(block) < self.gap_threshold
+                ]):
+                    block_structure.insert(prev_block._id, 0)
+                    selected_polygons.append(prev_block.polygon)
+                    page.structure.remove(prev_block._id)
+                else:
+                    break
+
+            for j, next_block in enumerate(page.children[i + 1:]):
+                if all([
+                    next_block.block_type in ["Caption", "Footnote"],
+                    next_block.minimum_gap(block) < self.gap_threshold
+                ]):
+                    block_structure.append(next_block._id)
+                    selected_polygons.append(next_block.polygon)
+                    page.structure.remove(next_block._id)
+                else:
+                    break
+
+            if len(block_structure) > 1:
+                # Create a merged block
+                new_block_cls = GROUP_BLOCK_REGISTRY[block.block_type + "Group"]
+                new_polygon = block.polygon.merge(selected_polygons)
+                group_block = page.add_block(new_block_cls, new_polygon)
+                group_block.structure = block_structure
+
+                # Update the structure of the page to reflect the new block
+                orig_block_idx = page.structure.index(block_structure[0])
+                page.structure[orig_block_idx] = group_block._id
+
+    def group_lists(self, page: PageGroup):
+        for i, block in enumerate(page.children):
+            if block.block_type not in ["ListItem"]:
+                continue
+            block_structure = [block._id]
+            selected_polygons = [block.polygon]
+
+            for j, next_block in enumerate(page.children[i + 1:]):
+                if all([
+                    next_block.block_type == "ListItem",
+                    next_block.minimum_gap(block) < self.gap_threshold
+                ]):
+                    block_structure.append(next_block._id)
+                    selected_polygons.append(next_block.polygon)
+                    page.structure.remove(next_block._id)
+                else:
+                    break
+
+            if len(block_structure) > 1:
+                new_polygon = block.polygon.merge(selected_polygons)
+                block = page.add_block(ListGroup, new_polygon)
+                block.structure = block_structure

marker/v2/converters/__init__.py

@@ -0,0 +1,16 @@
+from typing import List
+
+from pydantic import BaseModel
+
+
+class ConverterConfig(BaseModel):
+    filepath: str
+    page_range: List[int] | None = None
+
+
+class BaseConverter:
+    def __init__(self, config: ConverterConfig):
+        self.config = config
+
+    def __call__(self):
+        raise NotImplementedError

marker/v2/converters/pdf.py

@@ -0,0 +1,22 @@
+from surya.model.layout.model import load_model
+from surya.model.layout.processor import load_processor
+
+from marker.v2.builders.document import DocumentBuilder
+from marker.v2.builders.layout import LayoutBuilder
+from marker.v2.builders.structure import StructureBuilder
+from marker.v2.converters import BaseConverter
+from marker.v2.providers.pdf import PdfProvider
+from marker.v2.schema.config.pdf import PdfProviderConfig
+
+
+class PdfConverter(BaseConverter):
+    def __call__(self):
+        pdf_provider = PdfProvider(self.config.filepath, PdfProviderConfig())
+
+        layout_model = load_model()
+        layout_model.processor = load_processor()
+        layout_builder = LayoutBuilder(layout_model)
+
+        document = DocumentBuilder()(pdf_provider, layout_builder)
+        StructureBuilder()(document)
+

marker/v2/schema/__init__.py

@@ -16,8 +16,8 @@ class Block(BaseModel):
 
     @property
     def _id(self):
-        page_path = f"/page/{self.pnum}"
-        if self.block_num is not None:
-            return f"{page_path}/block/{self.block_num}"
+        page_path = f"/page/{self.page_id}"
+        if self.block_id is not None:
+            return f"{page_path}/block/{self.block_id}"
         else:
             return page_path

marker/v2/schema/groups/figure.py

@@ -2,4 +2,4 @@ from marker.v2.schema import Block
 
 
 class FigureGroup(Block):
-    pass
+    block_type: str = "FigureGroup"

marker/v2/schema/groups/list.py

@@ -2,4 +2,4 @@ from marker.v2.schema import Block
 
 
 class ListGroup(Block):
-    pass
+    block_type: str = "ListGroup"

marker/v2/schema/groups/picture.py

@@ -2,4 +2,4 @@ from marker.v2.schema import Block
 
 
 class PictureGroup(Block):
-    pass
+    block_type: str = "PictureGroup"

marker/v2/schema/groups/table.py

@@ -2,4 +2,4 @@ from marker.v2.schema import Block
 
 
 class TableGroup(Block):
-    pass
+    block_type: str = "TableGroup"

marker/v2/schema/polygon.py

@@ -83,13 +83,6 @@ class PolygonBox(BaseModel):
             corner[1] = max(min(corner[1], bounds[3]), bounds[1])
         self.polygon = new_corners
 
-    def merge(self, other):
-        x1 = min(self.bbox[0], other.bbox[0])
-        y1 = min(self.bbox[1], other.bbox[1])
-        x2 = max(self.bbox[2], other.bbox[2])
-        y2 = max(self.bbox[3], other.bbox[3])
-        self.polygon = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
-
     def overlap_x(self, other):
         return max(0, min(self.bbox[2], other.bbox[2]) - max(self.bbox[0], other.bbox[0]))
 
@@ -113,6 +106,14 @@ class PolygonBox(BaseModel):
         intersection = self.intersection_area(other, x_margin, y_margin)
         return intersection / self.area
 
+    def merge(self, others: List[PolygonBox]) -> PolygonBox:
+        corners = []
+        for i in range(len(self.polygon)):
+            min_x = min([self.polygon[i][0]] + [other.polygon[i][0] for other in others])
+            min_y = min([self.polygon[i][1]] + [other.polygon[i][1] for other in others])
+            corners.append([min_x, min_y])
+        return PolygonBox(polygon=corners)
+
     @classmethod
     def from_bbox(cls, bbox: List[float]):
         return cls(polygon=[[bbox[0], bbox[1]], [bbox[2], bbox[1]], [bbox[2], bbox[3]], [bbox[0], bbox[3]]])

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+testpaths=tests

tests/test_structure.py

@@ -0,0 +1,9 @@
+from marker.v2.builders.structure import StructureBuilder
+from tests.utils import setup_pdf_document
+
+
+def test_structure_builder():
+    document = setup_pdf_document('adversarial.pdf')
+    structure = StructureBuilder()
+    structure(document)
+    assert len(document.pages[0].structure) > 0

tests/utils.py

@@ -0,0 +1,28 @@
+import datasets
+from surya.model.layout.model import load_model
+from surya.model.layout.processor import load_processor
+import tempfile
+
+from marker.v2.builders.document import DocumentBuilder
+from marker.v2.builders.layout import LayoutBuilder
+from marker.v2.providers.pdf import PdfProvider
+from marker.v2.schema.config.pdf import PdfProviderConfig
+from marker.v2.schema.document import Document
+
+
+def setup_pdf_document(filename: str) -> Document:
+    dataset = datasets.load_dataset("datalab-to/pdfs", split="train")
+    idx = dataset['filename'].index(filename)
+
+    temp_pdf = tempfile.NamedTemporaryFile(suffix=".pdf")
+    temp_pdf.write(dataset['pdf'][idx])
+    temp_pdf.flush()
+
+    layout_model = load_model()
+    layout_model.processor = load_processor()
+
+    provider = PdfProvider(temp_pdf.name, PdfProviderConfig())
+    layout_builder = LayoutBuilder(layout_model)
+    builder = DocumentBuilder()
+    document = builder(provider, layout_builder)
+    return document