Skip to content

Commit b1e521f

Browse files
authored
align usage of content layer param (#326)
* align usage of content layer param Signed-off-by: Panos Vagenas <[email protected]> * parametrize content layers in visualizers Signed-off-by: Panos Vagenas <[email protected]> --------- Signed-off-by: Panos Vagenas <[email protected]>
1 parent c9af4d6 commit b1e521f

File tree

3 files changed

+15
-6
lines changed

3 files changed

+15
-6
lines changed

docling_core/transforms/visualizer/layout_visualizer.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ class Params(BaseModel):
4040
"""Layout visualization parameters."""
4141

4242
show_label: bool = True
43+
content_layers: set[ContentLayer] = {cl for cl in ContentLayer}
4344

4445
base_visualizer: Optional[BaseVisualizer] = None
4546
params: Params = Params()
@@ -141,9 +142,7 @@ def _draw_doc_layout(
141142
prev_image = None
142143
prev_page_nr = None
143144
for idx, (elem, _) in enumerate(
144-
doc.iterate_items(
145-
included_content_layers={ContentLayer.BODY, ContentLayer.FURNITURE}
146-
)
145+
doc.iterate_items(included_content_layers=self.params.content_layers)
147146
):
148147
if not isinstance(elem, DocItem):
149148
continue

docling_core/transforms/visualizer/reading_order_visualizer.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
from PIL import ImageDraw
77
from PIL.Image import Image
8+
from pydantic import BaseModel
89
from typing_extensions import override
910

1011
from docling_core.transforms.visualizer.base import BaseVisualizer
@@ -14,7 +15,16 @@
1415
class ReadingOrderVisualizer(BaseVisualizer):
1516
"""Reading order visualizer."""
1617

18+
class Params(BaseModel):
19+
"""Layout visualization parameters."""
20+
21+
show_label: bool = True
22+
content_layers: set[ContentLayer] = {
23+
cl for cl in ContentLayer if cl != ContentLayer.BACKGROUND
24+
}
25+
1726
base_visualizer: Optional[BaseVisualizer] = None
27+
params: Params = Params()
1828

1929
def _draw_arrow(
2030
self,
@@ -71,7 +81,7 @@ def _draw_doc_reading_order(
7181
my_images: dict[Optional[int], Image] = images or {}
7282
prev_page = None
7383
for elem, _ in doc.iterate_items(
74-
included_content_layers={ContentLayer.BODY, ContentLayer.FURNITURE},
84+
included_content_layers=self.params.content_layers,
7585
):
7686
if not isinstance(elem, DocItem):
7787
continue

docling_core/types/doc/document.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2873,7 +2873,7 @@ def print_element_tree(self):
28732873
self.iterate_items(
28742874
with_groups=True,
28752875
traverse_pictures=True,
2876-
included_content_layers={ContentLayer.BODY, ContentLayer.FURNITURE},
2876+
included_content_layers={cl for cl in ContentLayer},
28772877
)
28782878
):
28792879
if isinstance(item, GroupItem):
@@ -2897,7 +2897,7 @@ def export_to_element_tree(self) -> str:
28972897
self.iterate_items(
28982898
with_groups=True,
28992899
traverse_pictures=True,
2900-
included_content_layers={ContentLayer.BODY, ContentLayer.FURNITURE},
2900+
included_content_layers={cl for cl in ContentLayer},
29012901
)
29022902
):
29032903
if isinstance(item, GroupItem):

0 commit comments

Comments
 (0)