refactor: allow the usage of backends in the enrich models and generalize the interface (#742)

* fix get image with cropbox

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

* allow the usage of backends in the enrich models and generalize the interface

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

* move logic in BaseTextImageEnrichmentModel

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

* renaming

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

---------

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

docling/models/base_model.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import Any, Iterable
+from typing import Any, Generic, Iterable, Optional
 
-from docling_core.types.doc import DoclingDocument, NodeItem
+from docling_core.types.doc import DoclingDocument, NodeItem, TextItem
+from typing_extensions import TypeVar
 
-from docling.datamodel.base_models import Page
+from docling.datamodel.base_models import ItemAndImageEnrichmentElement, Page
 from docling.datamodel.document import ConversionResult
 
 
@@ -15,14 +16,54 @@ class BasePageModel(ABC):
         pass
 
 
-class BaseEnrichmentModel(ABC):
+EnrichElementT = TypeVar("EnrichElementT", default=NodeItem)
+
+
+class GenericEnrichmentModel(ABC, Generic[EnrichElementT]):
 
     @abstractmethod
     def is_processable(self, doc: DoclingDocument, element: NodeItem) -> bool:
         pass
 
     @abstractmethod
-    def __call__(
-        self, doc: DoclingDocument, element_batch: Iterable[NodeItem]
-    ) -> Iterable[Any]:
+    def prepare_element(
+        self, conv_res: ConversionResult, element: NodeItem
+    ) -> Optional[EnrichElementT]:
         pass
+
+    @abstractmethod
+    def __call__(
+        self, doc: DoclingDocument, element_batch: Iterable[EnrichElementT]
+    ) -> Iterable[NodeItem]:
+        pass
+
+
+class BaseEnrichmentModel(GenericEnrichmentModel[NodeItem]):
+
+    def prepare_element(
+        self, conv_res: ConversionResult, element: NodeItem
+    ) -> Optional[NodeItem]:
+        if self.is_processable(doc=conv_res.document, element=element):
+            return element
+        return None
+
+
+class BaseItemAndImageEnrichmentModel(
+    GenericEnrichmentModel[ItemAndImageEnrichmentElement]
+):
+
+    images_scale: float
+
+    def prepare_element(
+        self, conv_res: ConversionResult, element: NodeItem
+    ) -> Optional[ItemAndImageEnrichmentElement]:
+        if not self.is_processable(doc=conv_res.document, element=element):
+            return None
+
+        assert isinstance(element, TextItem)
+        element_prov = element.prov[0]
+        page_ix = element_prov.page_no - 1
+        cropped_image = conv_res.pages[page_ix].get_image(
+            scale=self.images_scale, cropbox=element_prov.bbox
+        )
+        return ItemAndImageEnrichmentElement(item=element, image=cropped_image)

docling/models/page_assemble_model.py

@@ -22,7 +22,7 @@ _log = logging.getLogger(__name__)
 
 
 class PageAssembleOptions(BaseModel):
-    keep_images: bool = False
+    pass
 
 
 class PageAssembleModel(BasePageModel):
@@ -174,11 +174,4 @@ class PageAssembleModel(BasePageModel):
                         elements=elements, headers=headers, body=body
                     )
 
-                    # Remove page images (can be disabled)
-                    if not self.options.keep_images:
-                        page._image_cache = {}
-
-                    # Unload backend
-                    page._backend.unload()
-
                 yield page