File tree Expand file tree Collapse file tree 3 files changed +15
-6
lines changed Expand file tree Collapse file tree 3 files changed +15
-6
lines changed Original file line number Diff line number Diff line change @@ -40,6 +40,7 @@ class Params(BaseModel):
40
40
"""Layout visualization parameters."""
41
41
42
42
show_label : bool = True
43
+ content_layers : set [ContentLayer ] = {cl for cl in ContentLayer }
43
44
44
45
base_visualizer : Optional [BaseVisualizer ] = None
45
46
params : Params = Params ()
@@ -141,9 +142,7 @@ def _draw_doc_layout(
141
142
prev_image = None
142
143
prev_page_nr = None
143
144
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 )
147
146
):
148
147
if not isinstance (elem , DocItem ):
149
148
continue
Original file line number Diff line number Diff line change 5
5
6
6
from PIL import ImageDraw
7
7
from PIL .Image import Image
8
+ from pydantic import BaseModel
8
9
from typing_extensions import override
9
10
10
11
from docling_core .transforms .visualizer .base import BaseVisualizer
14
15
class ReadingOrderVisualizer (BaseVisualizer ):
15
16
"""Reading order visualizer."""
16
17
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
+
17
26
base_visualizer : Optional [BaseVisualizer ] = None
27
+ params : Params = Params ()
18
28
19
29
def _draw_arrow (
20
30
self ,
@@ -71,7 +81,7 @@ def _draw_doc_reading_order(
71
81
my_images : dict [Optional [int ], Image ] = images or {}
72
82
prev_page = None
73
83
for elem , _ in doc .iterate_items (
74
- included_content_layers = { ContentLayer . BODY , ContentLayer . FURNITURE } ,
84
+ included_content_layers = self . params . content_layers ,
75
85
):
76
86
if not isinstance (elem , DocItem ):
77
87
continue
Original file line number Diff line number Diff line change @@ -2873,7 +2873,7 @@ def print_element_tree(self):
2873
2873
self .iterate_items (
2874
2874
with_groups = True ,
2875
2875
traverse_pictures = True ,
2876
- included_content_layers = {ContentLayer . BODY , ContentLayer . FURNITURE },
2876
+ included_content_layers = {cl for cl in ContentLayer },
2877
2877
)
2878
2878
):
2879
2879
if isinstance (item , GroupItem ):
@@ -2897,7 +2897,7 @@ def export_to_element_tree(self) -> str:
2897
2897
self .iterate_items (
2898
2898
with_groups = True ,
2899
2899
traverse_pictures = True ,
2900
- included_content_layers = {ContentLayer . BODY , ContentLayer . FURNITURE },
2900
+ included_content_layers = {cl for cl in ContentLayer },
2901
2901
)
2902
2902
):
2903
2903
if isinstance (item , GroupItem ):
You can’t perform that action at this time.
0 commit comments