|
@@ -106,7 +106,7 @@ def get_sub_regions_ocr_res(
|
|
|
return sub_regions_ocr_res
|
|
return sub_regions_ocr_res
|
|
|
|
|
|
|
|
|
|
|
|
|
-def calculate_iou(box1, box2):
|
|
|
|
|
|
|
+def _calculate_iou(box1, box2):
|
|
|
"""
|
|
"""
|
|
|
Calculate Intersection over Union (IoU) between two bounding boxes.
|
|
Calculate Intersection over Union (IoU) between two bounding boxes.
|
|
|
|
|
|
|
@@ -142,39 +142,69 @@ def calculate_iou(box1, box2):
|
|
|
return iou
|
|
return iou
|
|
|
|
|
|
|
|
|
|
|
|
|
-def _whether_overlaps_y_exceeds_threshold(bbox1, bbox2, overlap_ratio_threshold=0.6):
|
|
|
|
|
- _, y0_1, _, y1_1 = bbox1
|
|
|
|
|
- _, y0_2, _, y1_2 = bbox2
|
|
|
|
|
|
|
+def _whether_y_overlap_exceeds_threshold(bbox1, bbox2, overlap_ratio_threshold=0.6):
|
|
|
|
|
+ """
|
|
|
|
|
+ Determines whether the vertical overlap between two bounding boxes exceeds a given threshold.
|
|
|
|
|
+
|
|
|
|
|
+ Args:
|
|
|
|
|
+ bbox1 (tuple): The first bounding box defined as (left, top, right, bottom).
|
|
|
|
|
+ bbox2 (tuple): The second bounding box defined as (left, top, right, bottom).
|
|
|
|
|
+ overlap_ratio_threshold (float): The threshold ratio to determine if the overlap is significant.
|
|
|
|
|
+ Defaults to 0.6.
|
|
|
|
|
|
|
|
- overlap = max(0, min(y1_1, y1_2) - max(y0_1, y0_2))
|
|
|
|
|
- min_height = min(y1_1 - y0_1, y1_2 - y0_2)
|
|
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ bool: True if the vertical overlap divided by the minimum height of the two bounding boxes
|
|
|
|
|
+ exceeds the overlap_ratio_threshold, otherwise False.
|
|
|
|
|
+ """
|
|
|
|
|
+ _, y1_0, _, y1_1 = bbox1
|
|
|
|
|
+ _, y2_0, _, y2_1 = bbox2
|
|
|
|
|
+
|
|
|
|
|
+ overlap = max(0, min(y1_1, y2_1) - max(y1_0, y2_0))
|
|
|
|
|
+ min_height = min(y1_1 - y1_0, y2_1 - y2_0)
|
|
|
|
|
|
|
|
return (overlap / min_height) > overlap_ratio_threshold
|
|
return (overlap / min_height) > overlap_ratio_threshold
|
|
|
|
|
|
|
|
|
|
|
|
|
-def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_threshold=0.7):
|
|
|
|
|
- assert ocr_res["boxes"] and ocr_res["rec_texts"]
|
|
|
|
|
|
|
+def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_iou_threshold=0.7):
|
|
|
|
|
+ """
|
|
|
|
|
+ Sorts OCR results based on their spatial arrangement, grouping them into lines and blocks.
|
|
|
|
|
+
|
|
|
|
|
+ Args:
|
|
|
|
|
+ layout_bbox (tuple): A tuple representing the layout bounding box, defined as (left, top, right, bottom).
|
|
|
|
|
+ ocr_res (dict): A dictionary containing OCR results with the following keys:
|
|
|
|
|
+ - "boxes": A list of bounding boxes, each defined as [left, top, right, bottom].
|
|
|
|
|
+ - "rec_texts": A corresponding list of recognized text strings for each box.
|
|
|
|
|
+ line_height_iou_threshold (float): The threshold for determining whether two boxes belong to
|
|
|
|
|
+ the same line based on their vertical overlap. Defaults to 0.7.
|
|
|
|
|
+
|
|
|
|
|
+ Returns:
|
|
|
|
|
+ dict: A dictionary with the same structure as `ocr_res`, but with boxes and texts sorted
|
|
|
|
|
+ and grouped into lines and blocks.
|
|
|
|
|
+ """
|
|
|
|
|
+ assert (
|
|
|
|
|
+ ocr_res["boxes"] and ocr_res["rec_texts"]
|
|
|
|
|
+ ), "OCR results must contain 'boxes' and 'rec_texts'"
|
|
|
|
|
|
|
|
- # span->line->block
|
|
|
|
|
boxes = ocr_res["boxes"]
|
|
boxes = ocr_res["boxes"]
|
|
|
- rec_text = ocr_res["rec_texts"]
|
|
|
|
|
- x_min, x_max = layout_bbox[0], layout_bbox[2]
|
|
|
|
|
|
|
+ rec_texts = ocr_res["rec_texts"]
|
|
|
|
|
+
|
|
|
|
|
+ x_min, _, x_max, _ = layout_bbox
|
|
|
|
|
+
|
|
|
|
|
+ spans = list(zip(boxes, rec_texts))
|
|
|
|
|
|
|
|
- spans = list(zip(boxes, rec_text))
|
|
|
|
|
spans.sort(key=lambda span: span[0][1])
|
|
spans.sort(key=lambda span: span[0][1])
|
|
|
spans = [list(span) for span in spans]
|
|
spans = [list(span) for span in spans]
|
|
|
|
|
|
|
|
lines = []
|
|
lines = []
|
|
|
- first_span = spans[0]
|
|
|
|
|
- current_line = [first_span]
|
|
|
|
|
- current_y0, current_y1 = first_span[0][1], first_span[0][3]
|
|
|
|
|
|
|
+ current_line = [spans[0]]
|
|
|
|
|
+ current_y0, current_y1 = spans[0][0][1], spans[0][0][3]
|
|
|
|
|
|
|
|
for span in spans[1:]:
|
|
for span in spans[1:]:
|
|
|
y0, y1 = span[0][1], span[0][3]
|
|
y0, y1 = span[0][1], span[0][3]
|
|
|
- if _whether_overlaps_y_exceeds_threshold(
|
|
|
|
|
|
|
+ if _whether_y_overlap_exceeds_threshold(
|
|
|
(0, current_y0, 0, current_y1),
|
|
(0, current_y0, 0, current_y1),
|
|
|
(0, y0, 0, y1),
|
|
(0, y0, 0, y1),
|
|
|
- line_height_threshold,
|
|
|
|
|
|
|
+ line_height_iou_threshold,
|
|
|
):
|
|
):
|
|
|
current_line.append(span)
|
|
current_line.append(span)
|
|
|
current_y0 = min(current_y0, y0)
|
|
current_y0 = min(current_y0, y0)
|
|
@@ -191,13 +221,16 @@ def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_threshold=0.7):
|
|
|
line.sort(key=lambda span: span[0][0])
|
|
line.sort(key=lambda span: span[0][0])
|
|
|
first_span = line[0]
|
|
first_span = line[0]
|
|
|
end_span = line[-1]
|
|
end_span = line[-1]
|
|
|
- if first_span[0][0] - x_min > 20:
|
|
|
|
|
|
|
+
|
|
|
|
|
+ if first_span[0][0] - x_min > 15:
|
|
|
first_span[1] = "\n" + first_span[1]
|
|
first_span[1] = "\n" + first_span[1]
|
|
|
- if x_max - end_span[0][2] > 20:
|
|
|
|
|
|
|
+ if x_max - end_span[0][2] > 15:
|
|
|
end_span[1] = end_span[1] + "\n"
|
|
end_span[1] = end_span[1] + "\n"
|
|
|
|
|
|
|
|
|
|
+ # Flatten lines back into a single list for boxes and texts
|
|
|
ocr_res["boxes"] = [span[0] for line in lines for span in line]
|
|
ocr_res["boxes"] = [span[0] for line in lines for span in line]
|
|
|
ocr_res["rec_texts"] = [span[1] + " " for line in lines for span in line]
|
|
ocr_res["rec_texts"] = [span[1] + " " for line in lines for span in line]
|
|
|
|
|
+
|
|
|
return ocr_res
|
|
return ocr_res
|
|
|
|
|
|
|
|
|
|
|
|
@@ -241,7 +274,12 @@ def get_structure_res(
|
|
|
|
|
|
|
|
if label == "table":
|
|
if label == "table":
|
|
|
for i, table_res in enumerate(table_res_list):
|
|
for i, table_res in enumerate(table_res_list):
|
|
|
- if calculate_iou(layout_bbox, table_res["cell_box_list"][0]) > 0.5:
|
|
|
|
|
|
|
+ if (
|
|
|
|
|
+ _calculate_iou(
|
|
|
|
|
+ layout_bbox, table_res["table_ocr_pred"]["rec_boxes"][0]
|
|
|
|
|
+ )
|
|
|
|
|
+ > 0.5
|
|
|
|
|
+ ):
|
|
|
structure_boxes.append(
|
|
structure_boxes.append(
|
|
|
{
|
|
{
|
|
|
"label": label,
|
|
"label": label,
|
|
@@ -256,7 +294,7 @@ def get_structure_res(
|
|
|
else:
|
|
else:
|
|
|
overall_text_boxes = overall_ocr_res["rec_boxes"]
|
|
overall_text_boxes = overall_ocr_res["rec_boxes"]
|
|
|
for box_no in range(len(overall_text_boxes)):
|
|
for box_no in range(len(overall_text_boxes)):
|
|
|
- if calculate_iou(layout_bbox, overall_text_boxes[box_no]) > 0.5:
|
|
|
|
|
|
|
+ if _calculate_iou(layout_bbox, overall_text_boxes[box_no]) > 0.5:
|
|
|
rec_res["boxes"].append(overall_text_boxes[box_no])
|
|
rec_res["boxes"].append(overall_text_boxes[box_no])
|
|
|
rec_res["rec_texts"].append(
|
|
rec_res["rec_texts"].append(
|
|
|
overall_ocr_res["rec_texts"][box_no],
|
|
overall_ocr_res["rec_texts"][box_no],
|
|
@@ -306,7 +344,7 @@ def get_structure_res(
|
|
|
return structure_boxes
|
|
return structure_boxes
|
|
|
|
|
|
|
|
|
|
|
|
|
-def projection_by_bboxes(boxes: np.ndarray, axis: int) -> np.ndarray:
|
|
|
|
|
|
|
+def _projection_by_bboxes(boxes: np.ndarray, axis: int) -> np.ndarray:
|
|
|
"""
|
|
"""
|
|
|
Generate a 1D projection histogram from bounding boxes along a specified axis.
|
|
Generate a 1D projection histogram from bounding boxes along a specified axis.
|
|
|
|
|
|
|
@@ -328,7 +366,7 @@ def projection_by_bboxes(boxes: np.ndarray, axis: int) -> np.ndarray:
|
|
|
return projection
|
|
return projection
|
|
|
|
|
|
|
|
|
|
|
|
|
-def split_projection_profile(arr_values: np.ndarray, min_value: float, min_gap: float):
|
|
|
|
|
|
|
+def _split_projection_profile(arr_values: np.ndarray, min_value: float, min_gap: float):
|
|
|
"""
|
|
"""
|
|
|
Split the projection profile into segments based on specified thresholds.
|
|
Split the projection profile into segments based on specified thresholds.
|
|
|
|
|
|
|
@@ -363,7 +401,7 @@ def split_projection_profile(arr_values: np.ndarray, min_value: float, min_gap:
|
|
|
return segment_starts, segment_ends
|
|
return segment_starts, segment_ends
|
|
|
|
|
|
|
|
|
|
|
|
|
-def recursive_yx_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_gap=1):
|
|
|
|
|
|
|
+def _recursive_yx_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_gap=1):
|
|
|
"""
|
|
"""
|
|
|
Recursively project and segment bounding boxes, starting with Y-axis and followed by X-axis.
|
|
Recursively project and segment bounding boxes, starting with Y-axis and followed by X-axis.
|
|
|
|
|
|
|
@@ -380,8 +418,8 @@ def recursive_yx_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_
|
|
|
y_sorted_indices = np.array(indices)[y_sorted_indices]
|
|
y_sorted_indices = np.array(indices)[y_sorted_indices]
|
|
|
|
|
|
|
|
# Perform Y-axis projection
|
|
# Perform Y-axis projection
|
|
|
- y_projection = projection_by_bboxes(boxes=y_sorted_boxes, axis=1)
|
|
|
|
|
- y_intervals = split_projection_profile(y_projection, 0, 1)
|
|
|
|
|
|
|
+ y_projection = _projection_by_bboxes(boxes=y_sorted_boxes, axis=1)
|
|
|
|
|
+ y_intervals = _split_projection_profile(y_projection, 0, 1)
|
|
|
|
|
|
|
|
if not y_intervals:
|
|
if not y_intervals:
|
|
|
return
|
|
return
|
|
@@ -401,8 +439,8 @@ def recursive_yx_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_
|
|
|
x_sorted_indices_chunk = y_indices_chunk[x_sorted_indices]
|
|
x_sorted_indices_chunk = y_indices_chunk[x_sorted_indices]
|
|
|
|
|
|
|
|
# Perform X-axis projection
|
|
# Perform X-axis projection
|
|
|
- x_projection = projection_by_bboxes(boxes=x_sorted_boxes_chunk, axis=0)
|
|
|
|
|
- x_intervals = split_projection_profile(x_projection, 0, min_gap)
|
|
|
|
|
|
|
+ x_projection = _projection_by_bboxes(boxes=x_sorted_boxes_chunk, axis=0)
|
|
|
|
|
+ x_intervals = _split_projection_profile(x_projection, 0, min_gap)
|
|
|
|
|
|
|
|
if not x_intervals:
|
|
if not x_intervals:
|
|
|
continue
|
|
continue
|
|
@@ -417,14 +455,14 @@ def recursive_yx_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_
|
|
|
x_interval_indices = (x_start <= x_sorted_boxes_chunk[:, 0]) & (
|
|
x_interval_indices = (x_start <= x_sorted_boxes_chunk[:, 0]) & (
|
|
|
x_sorted_boxes_chunk[:, 0] < x_end
|
|
x_sorted_boxes_chunk[:, 0] < x_end
|
|
|
)
|
|
)
|
|
|
- recursive_yx_cut(
|
|
|
|
|
|
|
+ _recursive_yx_cut(
|
|
|
x_sorted_boxes_chunk[x_interval_indices],
|
|
x_sorted_boxes_chunk[x_interval_indices],
|
|
|
x_sorted_indices_chunk[x_interval_indices],
|
|
x_sorted_indices_chunk[x_interval_indices],
|
|
|
res,
|
|
res,
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
-def recursive_xy_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_gap=1):
|
|
|
|
|
|
|
+def _recursive_xy_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_gap=1):
|
|
|
"""
|
|
"""
|
|
|
Recursively performs X-axis projection followed by Y-axis projection to segment bounding boxes.
|
|
Recursively performs X-axis projection followed by Y-axis projection to segment bounding boxes.
|
|
|
|
|
|
|
@@ -442,8 +480,8 @@ def recursive_xy_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_
|
|
|
x_sorted_indices = np.array(indices)[x_sorted_indices]
|
|
x_sorted_indices = np.array(indices)[x_sorted_indices]
|
|
|
|
|
|
|
|
# Perform X-axis projection
|
|
# Perform X-axis projection
|
|
|
- x_projection = projection_by_bboxes(boxes=x_sorted_boxes, axis=0)
|
|
|
|
|
- x_intervals = split_projection_profile(x_projection, 0, 1)
|
|
|
|
|
|
|
+ x_projection = _projection_by_bboxes(boxes=x_sorted_boxes, axis=0)
|
|
|
|
|
+ x_intervals = _split_projection_profile(x_projection, 0, 1)
|
|
|
|
|
|
|
|
if not x_intervals:
|
|
if not x_intervals:
|
|
|
return
|
|
return
|
|
@@ -463,8 +501,8 @@ def recursive_xy_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_
|
|
|
y_sorted_indices_chunk = x_indices_chunk[y_sorted_indices]
|
|
y_sorted_indices_chunk = x_indices_chunk[y_sorted_indices]
|
|
|
|
|
|
|
|
# Perform Y-axis projection
|
|
# Perform Y-axis projection
|
|
|
- y_projection = projection_by_bboxes(boxes=y_sorted_boxes_chunk, axis=1)
|
|
|
|
|
- y_intervals = split_projection_profile(y_projection, 0, min_gap)
|
|
|
|
|
|
|
+ y_projection = _projection_by_bboxes(boxes=y_sorted_boxes_chunk, axis=1)
|
|
|
|
|
+ y_intervals = _split_projection_profile(y_projection, 0, min_gap)
|
|
|
|
|
|
|
|
if not y_intervals:
|
|
if not y_intervals:
|
|
|
continue
|
|
continue
|
|
@@ -479,7 +517,7 @@ def recursive_xy_cut(boxes: np.ndarray, indices: List[int], res: List[int], min_
|
|
|
y_interval_indices = (y_start <= y_sorted_boxes_chunk[:, 1]) & (
|
|
y_interval_indices = (y_start <= y_sorted_boxes_chunk[:, 1]) & (
|
|
|
y_sorted_boxes_chunk[:, 1] < y_end
|
|
y_sorted_boxes_chunk[:, 1] < y_end
|
|
|
)
|
|
)
|
|
|
- recursive_xy_cut(
|
|
|
|
|
|
|
+ _recursive_xy_cut(
|
|
|
y_sorted_boxes_chunk[y_interval_indices],
|
|
y_sorted_boxes_chunk[y_interval_indices],
|
|
|
y_sorted_indices_chunk[y_interval_indices],
|
|
y_sorted_indices_chunk[y_interval_indices],
|
|
|
res,
|
|
res,
|
|
@@ -490,7 +528,7 @@ def sort_by_xycut(block_bboxes, direction=0, min_gap=1):
|
|
|
block_bboxes = np.asarray(block_bboxes).astype(int)
|
|
block_bboxes = np.asarray(block_bboxes).astype(int)
|
|
|
res = []
|
|
res = []
|
|
|
if direction == 1:
|
|
if direction == 1:
|
|
|
- recursive_yx_cut(
|
|
|
|
|
|
|
+ _recursive_yx_cut(
|
|
|
block_bboxes,
|
|
block_bboxes,
|
|
|
np.arange(
|
|
np.arange(
|
|
|
len(block_bboxes),
|
|
len(block_bboxes),
|
|
@@ -499,7 +537,7 @@ def sort_by_xycut(block_bboxes, direction=0, min_gap=1):
|
|
|
min_gap,
|
|
min_gap,
|
|
|
)
|
|
)
|
|
|
else:
|
|
else:
|
|
|
- recursive_xy_cut(
|
|
|
|
|
|
|
+ _recursive_xy_cut(
|
|
|
block_bboxes,
|
|
block_bboxes,
|
|
|
np.arange(
|
|
np.arange(
|
|
|
len(block_bboxes),
|
|
len(block_bboxes),
|
|
@@ -1125,7 +1163,6 @@ def get_layout_ordering(data, no_mask_labels=[], already_sorted=False):
|
|
|
),
|
|
),
|
|
|
)
|
|
)
|
|
|
block_bboxes = np.array(block_bboxes)
|
|
block_bboxes = np.array(block_bboxes)
|
|
|
- print("sort by yxcut...")
|
|
|
|
|
sorted_indices = sort_by_xycut(
|
|
sorted_indices = sort_by_xycut(
|
|
|
block_bboxes,
|
|
block_bboxes,
|
|
|
direction=1,
|
|
direction=1,
|