feat: Factory and plugin-capability for Layout and Table models (#2637)

* feat: Scaffolding for layout and table model plugin factory

Signed-off-by: Christoph Auer <[email protected]>

* Add missing files

Signed-off-by: Christoph Auer <[email protected]>

* Add base options classes for layout and table

Signed-off-by: Christoph Auer <[email protected]>

---------

Signed-off-by: Christoph Auer <[email protected]>

Author: cau-git

docling/datamodel/pipeline_options.py

@@ -59,9 +59,14 @@ class TableFormerMode(str, Enum):
     ACCURATE = "accurate"
 
 
-class TableStructureOptions(BaseModel):
+class BaseTableStructureOptions(BaseOptions):
+    """Base options for table structure models."""
+
+
+class TableStructureOptions(BaseTableStructureOptions):
     """Options for the table structure."""
 
+    kind: ClassVar[str] = "docling_tableformer"
     do_cell_matching: bool = (
         True
         # True:  Matches predictions back to PDF cells. Can break table output if PDF cells
@@ -308,19 +313,25 @@ class VlmPipelineOptions(PaginatedPipelineOptions):
     )
 
 
-class LayoutOptions(BaseModel):
-    """Options for layout processing."""
+class BaseLayoutOptions(BaseOptions):
+    """Base options for layout models."""
 
-    create_orphan_clusters: bool = True  # Whether to create clusters for orphaned cells
     keep_empty_clusters: bool = (
         False  # Whether to keep clusters that contain no text cells
     )
-    model_spec: LayoutModelConfig = DOCLING_LAYOUT_HERON
     skip_cell_assignment: bool = (
         False  # Skip cell-to-cluster assignment for VLM-only processing
     )
 
 
+class LayoutOptions(BaseLayoutOptions):
+    """Options for layout processing."""
+
+    kind: ClassVar[str] = "docling_layout_default"
+    create_orphan_clusters: bool = True  # Whether to create clusters for orphaned cells
+    model_spec: LayoutModelConfig = DOCLING_LAYOUT_HERON
+
+
 class AsrPipelineOptions(PipelineOptions):
     asr_options: Union[InlineAsrOptions] = asr_model_specs.WHISPER_TINY
 

docling/models/base_layout_model.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+from typing import Type
+
+from docling.datamodel.base_models import LayoutPrediction, Page
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import BaseLayoutOptions
+from docling.models.base_model import BaseModelWithOptions, BasePageModel
+
+
+class BaseLayoutModel(BasePageModel, BaseModelWithOptions, ABC):
+    """Shared interface for layout models."""
+
+    @classmethod
+    @abstractmethod
+    def get_options_type(cls) -> Type[BaseLayoutOptions]:
+        """Return the options type supported by this layout model."""
+
+    @abstractmethod
+    def predict_layout(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[LayoutPrediction]:
+        """Produce layout predictions for the provided pages."""
+
+    def __call__(
+        self,
+        conv_res: ConversionResult,
+        page_batch: Iterable[Page],
+    ) -> Iterable[Page]:
+        pages = list(page_batch)
+        predictions = self.predict_layout(conv_res, pages)
+
+        for page, prediction in zip(pages, predictions):
+            page.predictions.layout = prediction
+            yield page

docling/models/base_table_model.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+from typing import Type
+
+from docling.datamodel.base_models import Page, TableStructurePrediction
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import BaseTableStructureOptions
+from docling.models.base_model import BaseModelWithOptions, BasePageModel
+
+
+class BaseTableStructureModel(BasePageModel, BaseModelWithOptions, ABC):
+    """Shared interface for table structure models."""
+
+    enabled: bool
+
+    @classmethod
+    @abstractmethod
+    def get_options_type(cls) -> Type[BaseTableStructureOptions]:
+        """Return the options type supported by this table model."""
+
+    @abstractmethod
+    def predict_tables(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[TableStructurePrediction]:
+        """Produce table structure predictions for the provided pages."""
+
+    def __call__(
+        self,
+        conv_res: ConversionResult,
+        page_batch: Iterable[Page],
+    ) -> Iterable[Page]:
+        if not getattr(self, "enabled", True):
+            yield from page_batch
+            return
+
+        pages = list(page_batch)
+        predictions = self.predict_tables(conv_res, pages)
+
+        for page, prediction in zip(pages, predictions):
+            page.predictions.tablestructure = prediction
+            yield page

docling/models/factories/__init__.py

@@ -1,10 +1,12 @@
 import logging
 from functools import lru_cache
 
+from docling.models.factories.layout_factory import LayoutFactory
 from docling.models.factories.ocr_factory import OcrFactory
 from docling.models.factories.picture_description_factory import (
     PictureDescriptionFactory,
 )
+from docling.models.factories.table_factory import TableStructureFactory
 
 logger = logging.getLogger(__name__)
 
@@ -25,3 +27,21 @@ def get_picture_description_factory(
     factory.load_from_plugins(allow_external_plugins=allow_external_plugins)
     logger.info("Registered picture descriptions: %r", factory.registered_kind)
     return factory
+
+
+@lru_cache
+def get_layout_factory(allow_external_plugins: bool = False) -> LayoutFactory:
+    factory = LayoutFactory()
+    factory.load_from_plugins(allow_external_plugins=allow_external_plugins)
+    logger.info("Registered layout engines: %r", factory.registered_kind)
+    return factory
+
+
+@lru_cache
+def get_table_structure_factory(
+    allow_external_plugins: bool = False,
+) -> TableStructureFactory:
+    factory = TableStructureFactory()
+    factory.load_from_plugins(allow_external_plugins=allow_external_plugins)
+    logger.info("Registered table structure engines: %r", factory.registered_kind)
+    return factory

docling/models/factories/layout_factory.py

@@ -0,0 +1,7 @@
+from docling.models.base_layout_model import BaseLayoutModel
+from docling.models.factories.base_factory import BaseFactory
+
+
+class LayoutFactory(BaseFactory[BaseLayoutModel]):
+    def __init__(self, *args, **kwargs):
+        super().__init__("layout_engines", *args, **kwargs)

docling/models/factories/table_factory.py

@@ -0,0 +1,7 @@
+from docling.models.base_table_model import BaseTableStructureModel
+from docling.models.factories.base_factory import BaseFactory
+
+
+class TableStructureFactory(BaseFactory[BaseTableStructureModel]):
+    def __init__(self, *args, **kwargs):
+        super().__init__("table_structure_engines", *args, **kwargs)

docling/models/layout_model.py

@@ -1,7 +1,7 @@
 import copy
 import logging
 import warnings
-from collections.abc import Iterable
+from collections.abc import Sequence
 from pathlib import Path
 from typing import List, Optional, Union
 
@@ -15,7 +15,7 @@
 from docling.datamodel.layout_model_specs import DOCLING_LAYOUT_V2, LayoutModelConfig
 from docling.datamodel.pipeline_options import LayoutOptions
 from docling.datamodel.settings import settings
-from docling.models.base_model import BasePageModel
+from docling.models.base_layout_model import BaseLayoutModel
 from docling.models.utils.hf_model_download import download_hf_model
 from docling.utils.accelerator_utils import decide_device
 from docling.utils.layout_postprocessor import LayoutPostprocessor
@@ -25,7 +25,7 @@
 _log = logging.getLogger(__name__)
 
 
-class LayoutModel(BasePageModel):
+class LayoutModel(BaseLayoutModel):
     TEXT_ELEM_LABELS = [
         DocItemLabel.TEXT,
         DocItemLabel.FOOTNOTE,
@@ -86,6 +86,10 @@ def __init__(
             num_threads=accelerator_options.num_threads,
         )
 
+    @classmethod
+    def get_options_type(cls) -> type[LayoutOptions]:
+        return LayoutOptions
+
     @staticmethod
     def download_models(
         local_dir: Optional[Path] = None,
@@ -145,11 +149,13 @@ def draw_clusters_and_cells_side_by_side(
             out_file = out_path / f"{mode_prefix}_layout_page_{page.page_no:05}.png"
             combined_image.save(str(out_file), format="png")
 
-    def __call__(
-        self, conv_res: ConversionResult, page_batch: Iterable[Page]
-    ) -> Iterable[Page]:
-        # Convert to list to allow multiple iterations
-        pages = list(page_batch)
+    def predict_layout(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[LayoutPrediction]:
+        # Convert to list to ensure predictable iteration
+        pages = list(pages)
 
         # Separate valid and invalid pages
         valid_pages = []
@@ -167,12 +173,6 @@ def __call__(
             valid_pages.append(page)
             valid_page_images.append(page_image)
 
-        _log.debug(f"{len(pages)=}")
-        if pages:
-            _log.debug(f"{pages[0].page_no}-{pages[-1].page_no}")
-        _log.debug(f"{len(valid_pages)=}")
-        _log.debug(f"{len(valid_page_images)=}")
-
         # Process all valid pages with batch prediction
         batch_predictions = []
         if valid_page_images:
@@ -182,11 +182,14 @@ def __call__(
                 )
 
         # Process each page with its predictions
+        layout_predictions: list[LayoutPrediction] = []
         valid_page_idx = 0
         for page in pages:
             assert page._backend is not None
             if not page._backend.is_valid():
-                yield page
+                existing_prediction = page.predictions.layout or LayoutPrediction()
+                page.predictions.layout = existing_prediction
+                layout_predictions.append(existing_prediction)
                 continue
 
             page_predictions = batch_predictions[valid_page_idx]
@@ -233,11 +236,14 @@ def __call__(
                     np.mean([c.confidence for c in processed_cells if c.from_ocr])
                 )
 
-            page.predictions.layout = LayoutPrediction(clusters=processed_clusters)
+            prediction = LayoutPrediction(clusters=processed_clusters)
+            page.predictions.layout = prediction
 
             if settings.debug.visualize_layout:
                 self.draw_clusters_and_cells_side_by_side(
                     conv_res, page, processed_clusters, mode_prefix="postprocessed"
                 )
 
-            yield page
+            layout_predictions.append(prediction)
+
+        return layout_predictions

docling/models/plugins/defaults.py

@@ -28,3 +28,23 @@ def picture_description():
             PictureDescriptionApiModel,
         ]
     }
+
+
+def layout_engines():
+    from docling.models.layout_model import LayoutModel
+
+    return {
+        "layout_engines": [
+            LayoutModel,
+        ]
+    }
+
+
+def table_structure_engines():
+    from docling.models.table_structure_model import TableStructureModel
+
+    return {
+        "table_structure_engines": [
+            TableStructureModel,
+        ]
+    }