Skip to content

Commit 8677d6e

Browse files
vagenasvagenas
authored
fix: fix multi-provenance item visualization (#277)
Signed-off-by: Panos Vagenas <[email protected]>
1 parent d05fe08 commit 8677d6e

File tree

2 files changed

+85
-80
lines changed

2 files changed

+85
-80
lines changed

‎docling_core/transforms/visualizer/layout_visualizer.py‎

Lines changed: 33 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -149,38 +149,41 @@ def _draw_doc_layout(
149149
continue
150150
if len(elem.prov) == 0:
151151
continue # Skip elements without provenances
152-
prov = elem.prov[0]
153-
page_nr = prov.page_no
154-
155-
if page_nr in my_images:
156-
image = my_images[page_nr]
157-
else:
158-
raise RuntimeError(f"Cannot visualize page-image for {page_nr}")
159-
160-
if prev_page_nr is None or page_nr > prev_page_nr: # new page begins
161-
# complete previous drawing
162-
if prev_page_nr is not None and prev_image and clusters:
163-
self._draw_clusters(
164-
image=prev_image,
165-
clusters=clusters,
166-
scale_x=prev_image.width / doc.pages[prev_page_nr].size.width,
167-
scale_y=prev_image.height / doc.pages[prev_page_nr].size.height,
168-
)
169-
clusters = []
170152

171-
tlo_bbox = prov.bbox.to_top_left_origin(
172-
page_height=doc.pages[prov.page_no].size.height
173-
)
174-
cluster = _TLCluster(
175-
id=idx,
176-
label=elem.label,
177-
brec=_TLBoundingRectangle.from_bounding_box(bbox=tlo_bbox),
178-
cells=[],
179-
)
180-
clusters.append(cluster)
153+
for prov in elem.prov:
154+
page_nr = prov.page_no
155+
156+
if page_nr in my_images:
157+
image = my_images[page_nr]
158+
else:
159+
raise RuntimeError(f"Cannot visualize page-image for {page_nr}")
160+
161+
if prev_page_nr is None or page_nr > prev_page_nr: # new page begins
162+
# complete previous drawing
163+
if prev_page_nr is not None and prev_image and clusters:
164+
self._draw_clusters(
165+
image=prev_image,
166+
clusters=clusters,
167+
scale_x=prev_image.width
168+
/ doc.pages[prev_page_nr].size.width,
169+
scale_y=prev_image.height
170+
/ doc.pages[prev_page_nr].size.height,
171+
)
172+
clusters = []
173+
174+
tlo_bbox = prov.bbox.to_top_left_origin(
175+
page_height=doc.pages[prov.page_no].size.height
176+
)
177+
cluster = _TLCluster(
178+
id=idx,
179+
label=elem.label,
180+
brec=_TLBoundingRectangle.from_bounding_box(bbox=tlo_bbox),
181+
cells=[],
182+
)
183+
clusters.append(cluster)
181184

182-
prev_page_nr = page_nr
183-
prev_image = image
185+
prev_page_nr = page_nr
186+
prev_image = image
184187

185188
# complete last drawing
186189
if prev_page_nr is not None and prev_image and clusters:

‎docling_core/transforms/visualizer/reading_order_visualizer.py‎

Lines changed: 52 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -77,57 +77,59 @@ def _draw_doc_reading_order(
7777
continue
7878
if len(elem.prov) == 0:
7979
continue # Skip elements without provenances
80-
prov = elem.prov[0]
81-
page_no = prov.page_no
82-
image = my_images.get(page_no)
83-
84-
if image is None or prev_page is None or page_no > prev_page:
85-
# new page begins
86-
prev_page = page_no
87-
x0 = y0 = None
88-
89-
if image is None:
90-
page_image = doc.pages[page_no].image
91-
if page_image is None or (pil_img := page_image.pil_image) is None:
92-
raise RuntimeError("Cannot visualize document without images")
93-
else:
94-
image = deepcopy(pil_img)
95-
my_images[page_no] = image
96-
draw = ImageDraw.Draw(image)
97-
98-
# if prov.page_no not in true_doc.pages or prov.page_no != 1:
99-
# logging.error(f"{prov.page_no} not in true_doc.pages -> skipping! ")
100-
# continue
101-
102-
tlo_bbox = prov.bbox.to_top_left_origin(
103-
page_height=doc.pages[prov.page_no].size.height
104-
)
105-
ro_bbox = tlo_bbox.normalized(doc.pages[prov.page_no].size)
106-
ro_bbox.l = round(ro_bbox.l * image.width) # noqa: E741
107-
ro_bbox.r = round(ro_bbox.r * image.width)
108-
ro_bbox.t = round(ro_bbox.t * image.height)
109-
ro_bbox.b = round(ro_bbox.b * image.height)
110-
111-
if ro_bbox.b > ro_bbox.t:
112-
ro_bbox.b, ro_bbox.t = ro_bbox.t, ro_bbox.b
113-
114-
if x0 is None and y0 is None:
115-
x0 = (ro_bbox.l + ro_bbox.r) / 2.0
116-
y0 = (ro_bbox.b + ro_bbox.t) / 2.0
117-
else:
118-
assert x0 is not None
119-
assert y0 is not None
120-
121-
x1 = (ro_bbox.l + ro_bbox.r) / 2.0
122-
y1 = (ro_bbox.b + ro_bbox.t) / 2.0
123-
124-
draw = self._draw_arrow(
125-
draw=draw,
126-
arrow_coords=(x0, y0, x1, y1),
127-
line_width=2,
128-
color="red",
80+
81+
for prov in elem.prov:
82+
page_no = prov.page_no
83+
image = my_images.get(page_no)
84+
85+
if image is None or prev_page is None or page_no > prev_page:
86+
# new page begins
87+
prev_page = page_no
88+
x0 = y0 = None
89+
90+
if image is None:
91+
page_image = doc.pages[page_no].image
92+
if (
93+
page_image is None
94+
or (pil_img := page_image.pil_image) is None
95+
):
96+
raise RuntimeError(
97+
"Cannot visualize document without images"
98+
)
99+
else:
100+
image = deepcopy(pil_img)
101+
my_images[page_no] = image
102+
draw = ImageDraw.Draw(image)
103+
104+
tlo_bbox = prov.bbox.to_top_left_origin(
105+
page_height=doc.pages[prov.page_no].size.height
129106
)
130-
x0, y0 = x1, y1
107+
ro_bbox = tlo_bbox.normalized(doc.pages[prov.page_no].size)
108+
ro_bbox.l = round(ro_bbox.l * image.width) # noqa: E741
109+
ro_bbox.r = round(ro_bbox.r * image.width)
110+
ro_bbox.t = round(ro_bbox.t * image.height)
111+
ro_bbox.b = round(ro_bbox.b * image.height)
112+
113+
if ro_bbox.b > ro_bbox.t:
114+
ro_bbox.b, ro_bbox.t = ro_bbox.t, ro_bbox.b
115+
116+
if x0 is None and y0 is None:
117+
x0 = (ro_bbox.l + ro_bbox.r) / 2.0
118+
y0 = (ro_bbox.b + ro_bbox.t) / 2.0
119+
else:
120+
assert x0 is not None
121+
assert y0 is not None
122+
123+
x1 = (ro_bbox.l + ro_bbox.r) / 2.0
124+
y1 = (ro_bbox.b + ro_bbox.t) / 2.0
125+
126+
draw = self._draw_arrow(
127+
draw=draw,
128+
arrow_coords=(x0, y0, x1, y1),
129+
line_width=2,
130+
color="red",
131+
)
132+
x0, y0 = x1, y1
131133
return my_images
132134

133135
@override

0 commit comments

Comments
 (0)