Skip to content

Commit 5cb0859

Browse files
committed
fix(ocr): fix layout debug
Signed-off-by: Clément Doumouro <[email protected]>
1 parent cfd3635 commit 5cb0859

File tree

1 file changed

+30
-8
lines changed

1 file changed

+30
-8
lines changed

docling/models/layout_model.py

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import copy
22
import logging
33
import warnings
4+
from copy import deepcopy
45
from collections.abc import Iterable
56
from pathlib import Path
67
from typing import List, Optional, Union
@@ -103,17 +104,27 @@ def download_models(
103104
)
104105

105106
def draw_clusters_and_cells_side_by_side(
106-
self, conv_res, page, clusters, mode_prefix: str, show: bool = False
107+
self,
108+
conv_res,
109+
page,
110+
page_orientation: int,
111+
clusters,
112+
mode_prefix: str,
113+
show: bool = False,
107114
):
108115
"""
109116
Draws a page image side by side with clusters filtered into two categories:
110117
- Left: Clusters excluding FORM, KEY_VALUE_REGION, and PICTURE.
111118
- Right: Clusters including FORM, KEY_VALUE_REGION, and PICTURE.
112119
Includes label names and confidence scores for each cluster.
113120
"""
114-
scale_x = page.image.width / page.size.width
115-
scale_y = page.image.height / page.size.height
116-
121+
page_image = deepcopy(page.image)
122+
scale_x = page_image.width / page.size.width
123+
scale_y = page_image.height / page.size.height
124+
if page_orientation:
125+
page_image = page_image.rotate(-page_orientation, expand=True)
126+
if abs(page_orientation) in [90, 270]:
127+
scale_x, scale_y = scale_y, scale_x
117128
# Filter clusters for left and right images
118129
exclude_labels = {
119130
DocItemLabel.FORM,
@@ -123,12 +134,15 @@ def draw_clusters_and_cells_side_by_side(
123134
left_clusters = [c for c in clusters if c.label not in exclude_labels]
124135
right_clusters = [c for c in clusters if c.label in exclude_labels]
125136
# Create a deep copy of the original image for both sides
126-
left_image = copy.deepcopy(page.image)
127-
right_image = copy.deepcopy(page.image)
137+
left_image = page_image
138+
right_image = copy.deepcopy(left_image)
128139

129140
# Draw clusters on both images
130141
draw_clusters(left_image, left_clusters, scale_x, scale_y)
131142
draw_clusters(right_image, right_clusters, scale_x, scale_y)
143+
if page_orientation:
144+
left_image = left_image.rotate(page_orientation, expand=True)
145+
right_image = right_image.rotate(page_orientation, expand=True)
132146
# Combine the images side by side
133147
combined_width = left_image.width * 2
134148
combined_height = left_image.height
@@ -207,7 +221,11 @@ def __call__(
207221

208222
if settings.debug.visualize_raw_layout:
209223
self.draw_clusters_and_cells_side_by_side(
210-
conv_res, page, clusters, mode_prefix="raw"
224+
conv_res,
225+
page,
226+
page_orientation,
227+
clusters,
228+
mode_prefix="raw",
211229
)
212230

213231
# Apply postprocessing
@@ -236,7 +254,11 @@ def __call__(
236254

237255
if settings.debug.visualize_layout:
238256
self.draw_clusters_and_cells_side_by_side(
239-
conv_res, page, processed_clusters, mode_prefix="postprocessed"
257+
conv_res,
258+
page,
259+
page_orientation,
260+
processed_clusters,
261+
mode_prefix="postprocessed",
240262
)
241263

242264
yield page

0 commit comments

Comments
 (0)