feat: provide visualizer option in HTML split view (#294)

vagenas · web-flow · commit 6a7eb537eb6a · 2025-05-16T11:22:12.000+02:00
* feat: provide visualizer option in HTML split view

Signed-off-by: Panos Vagenas &lt;pva@zurich.ibm.com&gt;

* loosen test

Signed-off-by: Panos Vagenas &lt;pva@zurich.ibm.com&gt;

---------

Signed-off-by: Panos Vagenas &lt;pva@zurich.ibm.com&gt;
diff --git a/docling_core/transforms/serializer/common.py b/docling_core/transforms/serializer/common.py
@@ -169,7 +169,7 @@ class CommonParams(BaseModel):
 
     def merge_with_patch(self, patch: dict[str, Any]) -> Self:
         """Create an instance by merging the provided patch dict on top of self."""
-        res = self.model_validate({**self.model_dump(), **patch})
+        res = self.model_copy(update=patch)
         return res
 
 
@@ -260,10 +260,10 @@ def serialize_doc(
         """Serialize a document out of its pages."""
         ...
 
-    def _serialize_body(self) -> SerializationResult:
+    def _serialize_body(self, **kwargs) -> SerializationResult:
         """Serialize the document body."""
         subparts = self.get_parts()
-        res = self.serialize_doc(parts=subparts)
+        res = self.serialize_doc(parts=subparts, **kwargs)
         return res
 
     @override
@@ -278,12 +278,12 @@ def serialize(
     ) -> SerializationResult:
         """Serialize a given node."""
         my_visited: set[str] = visited if visited is not None else set()
-        my_kwargs = self.params.merge_with_patch(patch=kwargs).model_dump()
+        my_kwargs = {**self.params.model_dump(), **kwargs}
         empty_res = create_ser_result()
         if item is None or item == self.doc.body:
             if self.doc.body.self_ref not in my_visited:
                 my_visited.add(self.doc.body.self_ref)
-                return self._serialize_body()
+                return self._serialize_body(**my_kwargs)
             else:
                 return empty_res
 
diff --git a/docling_core/transforms/serializer/html.py b/docling_core/transforms/serializer/html.py
@@ -16,6 +16,7 @@
 from xml.sax.saxutils import unescape
 
 import latex2mathml.converter
+from PIL.Image import Image
 from pydantic import AnyUrl, BaseModel
 from typing_extensions import override
 
@@ -40,6 +41,7 @@
     _get_css_for_single_column,
     _get_css_for_split_page,
 )
+from docling_core.transforms.visualizer.base import BaseVisualizer
 from docling_core.types.doc.base import ImageRefMode
 from docling_core.types.doc.document import (
     CodeItem,
@@ -821,9 +823,22 @@ def serialize_hyperlink(
     def serialize_doc(
         self,
         parts: list[SerializationResult],
+        visualizer: Optional[BaseVisualizer] = None,
         **kwargs: Any,
     ) -> SerializationResult:
         """Serialize a document out of its pages."""
+
+        def _serialize_page_img(page_img: Image):
+            buffered = BytesIO()
+            page_img.save(buffered, format="PNG")  # Save the image to the byte stream
+            img_bytes = buffered.getvalue()  # Get the byte data
+
+            # Encode to Base64 and decode to string
+            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
+            img_text = f'<img src="data:image/png;base64,{img_base64}">'
+
+            return f"<figure>{img_text}</figure>"
+
         # Create HTML structure
         html_parts = [
             "<!DOCTYPE html>",
@@ -853,19 +868,26 @@ def serialize_doc(
             html_parts.append("<table>")
             html_parts.append("<tbody>")
 
+            vized_pages_dict: dict[Optional[int], Image] = {}
+            if visualizer:
+                vized_pages_dict = visualizer.get_visualization(doc=self.doc)
+
             for page_no, page in pages.items():
 
                 if isinstance(page_no, int):
                     if applicable_pages is not None and page_no not in applicable_pages:
                         continue
                     page_img = self.doc.pages[page_no].image
+                    vized_page = vized_pages_dict.get(page_no)
 
                     html_parts.append("<tr>")
 
                     html_parts.append("<td>")
 
+                    if vized_page:
+                        html_parts.append(_serialize_page_img(page_img=vized_page))
                     # short-cut: we already have the image in base64
-                    if (
+                    elif (
                         (page_img is not None)
                         and isinstance(page_img, ImageRef)
                         and isinstance(page_img.uri, AnyUrl)
@@ -875,18 +897,7 @@ def serialize_doc(
                         html_parts.append(f"<figure>{img_text}</figure>")
 
                     elif (page_img is not None) and (page_img._pil is not None):
-
-                        buffered = BytesIO()
-                        page_img._pil.save(
-                            buffered, format="PNG"
-                        )  # Save the image to the byte stream
-                        img_bytes = buffered.getvalue()  # Get the byte data
-
-                        # Encode to Base64 and decode to string
-                        img_base64 = base64.b64encode(img_bytes).decode("utf-8")
-                        img_text = f'<img src="data:image/png;base64,{img_base64}">'
-
-                        html_parts.append(f"<figure>{img_text}</figure>")
+                        html_parts.append(_serialize_page_img(page_img=page_img._pil))
                     else:
                         html_parts.append("<figure>no page-image found</figure>")
 
diff --git a/test/test_serialization.py b/test/test_serialization.py
@@ -12,6 +12,7 @@
     MarkdownDocSerializer,
     MarkdownParams,
 )
+from docling_core.transforms.visualizer.layout_visualizer import LayoutVisualizer
 from docling_core.types.doc.base import ImageRefMode
 from docling_core.types.doc.document import DoclingDocument
 from docling_core.types.doc.labels import DocItemLabel
@@ -183,6 +184,32 @@ def test_html_split_page_p2():
     verify(exp_file=src.parent / f"{src.stem}_split_p2.gt.html", actual=actual)
 
 
+def test_html_split_page_p2_with_visualizer():
+    src = Path("./test/data/doc/2408.09869v3_enriched.json")
+    doc = DoclingDocument.load_from_json(src)
+
+    ser = HTMLDocSerializer(
+        doc=doc,
+        params=HTMLParams(
+            image_mode=ImageRefMode.EMBEDDED,
+            output_style=HTMLOutputStyle.SPLIT_PAGE,
+            pages={2},
+        ),
+    )
+    ser_res = ser.serialize(
+        visualizer=LayoutVisualizer(),
+    )
+    actual = ser_res.text
+
+    # pinning the result with visualizer appeared flaky, so at least ensure it contains
+    # a figure (for the page) and that it is different than without visualizer:
+    assert '<figure><img src="data:image/png;base64' in actual
+    file_without_viz = src.parent / f"{src.stem}_split_p2.gt.html"
+    with open(file_without_viz) as f:
+        data_without_viz = f.read()
+    assert actual.strip() != data_without_viz.strip()
+
+
 def test_html_split_page_no_page_breaks():
     src = Path("./test/data/doc/2408.09869_p1.json")
     doc = DoclingDocument.load_from_json(src)
