parametrize content layers in visualizers

Signed-off-by: Panos Vagenas <[email protected]>

Author: vagenas

docling_core/transforms/visualizer/layout_visualizer.py

@@ -40,6 +40,7 @@ class Params(BaseModel):
         """Layout visualization parameters."""
 
         show_label: bool = True
+        content_layers: set[ContentLayer] = {cl for cl in ContentLayer}
 
     base_visualizer: Optional[BaseVisualizer] = None
     params: Params = Params()
@@ -141,7 +142,7 @@ def _draw_doc_layout(
         prev_image = None
         prev_page_nr = None
         for idx, (elem, _) in enumerate(
-            doc.iterate_items(included_content_layers={cl for cl in ContentLayer})
+            doc.iterate_items(included_content_layers=self.params.content_layers)
         ):
             if not isinstance(elem, DocItem):
                 continue

docling_core/transforms/visualizer/reading_order_visualizer.py

@@ -5,6 +5,7 @@
 
 from PIL import ImageDraw
 from PIL.Image import Image
+from pydantic import BaseModel
 from typing_extensions import override
 
 from docling_core.transforms.visualizer.base import BaseVisualizer
@@ -14,7 +15,16 @@
 class ReadingOrderVisualizer(BaseVisualizer):
     """Reading order visualizer."""
 
+    class Params(BaseModel):
+        """Layout visualization parameters."""
+
+        show_label: bool = True
+        content_layers: set[ContentLayer] = {
+            cl for cl in ContentLayer if cl != ContentLayer.BACKGROUND
+        }
+
     base_visualizer: Optional[BaseVisualizer] = None
+    params: Params = Params()
 
     def _draw_arrow(
         self,
@@ -71,7 +81,7 @@ def _draw_doc_reading_order(
         my_images: dict[Optional[int], Image] = images or {}
         prev_page = None
         for elem, _ in doc.iterate_items(
-            included_content_layers={cl for cl in ContentLayer},
+            included_content_layers=self.params.content_layers,
         ):
             if not isinstance(elem, DocItem):
                 continue