Skip to content

Commit

Permalink
Merge pull request #78 from icecraft/debug/txt_pipeline
Browse files Browse the repository at this point in the history
  • Loading branch information
@myhloli
myhloli authored Apr 28, 2024
2 parents 683fa63 + 6a3d1f2 commit a1a1d17
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 30 deletions.
67 changes: 39 additions & 28 deletions magic_pdf/model/magic_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,16 @@
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
from magic_pdf.libs.math import float_gt
from magic_pdf.libs.boxbase import _is_in, bbox_relative_pos, bbox_distance
from magic_pdf.libs.boxbase import (
_is_in,
bbox_relative_pos,
bbox_distance,
_is_part_overlap,
calculate_overlap_area_in_bbox1_area_ratio,
)
from magic_pdf.libs.ModelBlockTypeEnum import ModelBlockTypeEnum

CAPATION_OVERLAP_AREA_RATIO = 0.6

class MagicModel:
"""
Expand Down Expand Up @@ -74,13 +81,13 @@ def __reduct_overlap(self, bboxes):
return [bboxes[i] for i in range(N) if keep[i]]

def __tie_up_category_by_distance(
self, page_no, subject_category_id, object_category_id
self, page_no, subject_category_id, object_category_id
):
"""
假定每个 subject 最多有一个 object (可以有多个相邻的 object 合并为单个 object)，每个 object 只能属于一个 subject
"""
ret = []
MAX_DIS_OF_POINT = 10 ** 9 + 7
MAX_DIS_OF_POINT = 10**9 + 7

subjects = self.__reduct_overlap(
list(
Expand Down Expand Up @@ -123,8 +130,8 @@ def __tie_up_category_by_distance(
for i in range(N):
for j in range(i):
if (
all_bboxes[i]["category_id"] == subject_category_id
and all_bboxes[j]["category_id"] == subject_category_id
all_bboxes[i]["category_id"] == subject_category_id
and all_bboxes[j]["category_id"] == subject_category_id
):
continue

Expand Down Expand Up @@ -154,9 +161,9 @@ def __tie_up_category_by_distance(
if pos_flag_count > 1:
continue
if (
all_bboxes[j]["category_id"] != object_category_id
or j in used
or dis[i][j] == MAX_DIS_OF_POINT
all_bboxes[j]["category_id"] != object_category_id
or j in used
or dis[i][j] == MAX_DIS_OF_POINT
):
continue
arr.append((dis[i][j], j))
Expand Down Expand Up @@ -185,10 +192,10 @@ def __tie_up_category_by_distance(
continue

if (
all_bboxes[k]["category_id"] != object_category_id
or k in used
or k in seen
or dis[j][k] == MAX_DIS_OF_POINT
all_bboxes[k]["category_id"] != object_category_id
or k in used
or k in seen
or dis[j][k] == MAX_DIS_OF_POINT
):
continue
is_nearest = True
Expand Down Expand Up @@ -238,7 +245,7 @@ def __tie_up_category_by_distance(
for bbox in caption_poses:
embed_arr = []
for idx in seen:
if _is_in(all_bboxes[idx]["bbox"], bbox):
if calculate_overlap_area_in_bbox1_area_ratio(all_bboxes[idx]["bbox"], bbox) > CAPATION_OVERLAP_AREA_RATIO:
embed_arr.append(idx)

if len(embed_arr) > 0:
Expand All @@ -258,7 +265,7 @@ def __tie_up_category_by_distance(
caption_bbox = caption_poses[max_area_idx]

for j in seen:
if _is_in(all_bboxes[j]["bbox"], caption_bbox):
if calculate_overlap_area_in_bbox1_area_ratio(all_bboxes[j]["bbox"], caption_bbox) > CAPATION_OVERLAP_AREA_RATIO:
used.add(j)
subject_object_relation_map[i].append(j)

Expand Down Expand Up @@ -312,8 +319,8 @@ def __tie_up_category_by_distance(
candidates = []
for j in range(N):
if (
all_bboxes[j]["category_id"] != subject_category_id
or j in with_caption_subject
all_bboxes[j]["category_id"] != subject_category_id
or j in with_caption_subject
):
continue
candidates.append((dis[i][j], j))
Expand All @@ -335,7 +342,7 @@ def get_imgs(self, page_no: int): # @许瑞
]

def get_tables(
self, page_no: int
self, page_no: int
) -> list: # 3个坐标， caption, table主体，table-note
with_captions, _ = self.__tie_up_category_by_distance(page_no, 5, 6)
with_footnotes, _ = self.__tie_up_category_by_distance(page_no, 5, 7)
Expand All @@ -358,9 +365,15 @@ def get_tables(
return ret

def get_equations(self, page_no: int) -> list: # 有坐标，也有字
inline_equations = self.__get_blocks_by_type(ModelBlockTypeEnum.EMBEDDING.value, page_no, ["latex"])
interline_equations = self.__get_blocks_by_type(ModelBlockTypeEnum.ISOLATED.value, page_no, ["latex"])
interline_equations_blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.ISOLATE_FORMULA.value, page_no)
inline_equations = self.__get_blocks_by_type(
ModelBlockTypeEnum.EMBEDDING.value, page_no, ["latex"]
)
interline_equations = self.__get_blocks_by_type(
ModelBlockTypeEnum.ISOLATED.value, page_no, ["latex"]
)
interline_equations_blocks = self.__get_blocks_by_type(
ModelBlockTypeEnum.ISOLATE_FORMULA.value, page_no
)
return inline_equations, interline_equations, interline_equations_blocks

def get_discarded(self, page_no: int) -> list: # 自研模型，只有坐标
Expand All @@ -382,7 +395,7 @@ def get_ocr_text(self, page_no: int) -> list: # paddle 搞的，有字也有坐
for layout_det in layout_dets:
if layout_det["category_id"] == "15":
span = {
"bbox": layout_det['bbox'],
"bbox": layout_det["bbox"],
"content": layout_det["text"],
}
text_spans.append(span)
Expand All @@ -402,9 +415,7 @@ def get_all_spans(self, page_no: int) -> list:
for layout_det in layout_dets:
category_id = layout_det["category_id"]
if category_id in allow_category_id_list:
span = {
"bbox": layout_det['bbox']
}
span = {"bbox": layout_det["bbox"]}
if category_id == 3:
span["type"] = ContentType.Image
elif category_id == 5:
Expand All @@ -429,7 +440,9 @@ def get_page_size(self, page_no: int): # 获取页面宽高
page_h = page.rect.height
return page_w, page_h

def __get_blocks_by_type(self, type: int, page_no: int, extra_col: list[str] = []) -> list:
def __get_blocks_by_type(
self, type: int, page_no: int, extra_col: list[str] = []
) -> list:
blocks = []
for page_dict in self.__model_list:
layout_dets = page_dict.get("layout_dets", [])
Expand All @@ -442,9 +455,7 @@ def __get_blocks_by_type(self, type: int, page_no: int, extra_col: list[str] = [
bbox = item.get("bbox", None)

if category_id == type:
block = {
"bbox": bbox
}
block = {"bbox": bbox}
for col in extra_col:
block[col] = item.get(col, None)
blocks.append(block)
Expand Down
16 changes: 14 additions & 2 deletions magic_pdf/pre_proc/remove_bbox_overlap.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,21 @@

def _remove_overlap_between_bbox(spans):
res = []
for v in spans:

keeps = [True] * len(spans)
for i in range(len(spans)):
for j in range(len(spans)):
if i == j:
continue
if _is_in(spans[i]["bbox"], spans[j]["bbox"]):
keeps[i] = False

for idx, v in enumerate(spans):
if not keeps[idx]:
continue

for i in range(len(res)):
if _is_in(res[i]["bbox"], v["bbox"]) or _is_in(v["bbox"], res[i]["bbox"]):
if _is_in(v["bbox"], res[i]["bbox"]):
continue
if _is_in_or_part_overlap(res[i]["bbox"], v["bbox"]):
ix0, iy0, ix1, iy1 = res[i]["bbox"]
Expand Down

0 comments on commit a1a1d17

Please sign in to comment.