Skip to content

P
pdf-miner
Commit 7da3b545 authored by liukaiwen's avatar liukaiwen
Browse files
Options

更新了magic_model

parent f5b07d82
Hide whitespace changes
Inline Side-by-side
Showing with 0 additions and 418 deletions
magic_pdf/model/magic_model.py
View file @ 7da3b545
...@@ -440,424 +440,6 @@ class MagicModel: ...@@ -440,424 +440,6 @@ class MagicModel:
blocks.append(block) blocks.append(block)
return blocks return blocks
if __name__ == "__main__":
drw = DiskReaderWriter(r"D:/project/20231108code-clean")
if 0:
pdf_file_path = r"linshixuqiu\19983-00.pdf"
model_file_path = r"linshixuqiu\19983-00_new.json"
pdf_bytes = drw.read(pdf_file_path, AbsReaderWriter.MODE_BIN)
model_json_txt = drw.read(model_file_path, AbsReaderWriter.MODE_TXT)
model_list = json.loads(model_json_txt)
write_path = r"D:\project\20231108code-clean\linshixuqiu\19983-00"
img_bucket_path = "imgs"
img_writer = DiskReaderWriter(join_path(write_path, img_bucket_path))
pdf_docs = fitz.open("pdf", pdf_bytes)
magic_model = MagicModel(model_list, pdf_docs)
if 1:
model_list = json.loads(
drw.read("/opt/data/pdf/20240418/j.chroma.2009.03.042.json")
)
pdf_bytes = drw.read(
"/opt/data/pdf/20240418/j.chroma.2009.03.042.pdf", AbsReaderWriter.MODE_BIN
)
pdf_docs = fitz.open("pdf", pdf_bytes)
magic_model = MagicModel(model_list, pdf_docs)
for i in range(7):
print(magic_model.get_imgs(i))
def __reduct_overlap(self, bboxes):
N = len(bboxes)
keep = [True] * N
for i in range(N):
for j in range(N):
if i == j:
continue
if _is_in(bboxes[i], bboxes[j]):
keep[i] = False
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
):
"""
假定每个 subject 最多有一个 object (可以有多个相邻的 object 合并为单个 object)，每个 object 只能属于一个 subject
"""
ret = []
MAX_DIS_OF_POINT = 10**9 + 7
subjects = self.__reduct_overlap(
list(
map(
lambda x: x["bbox"],
filter(
lambda x: x["category_id"] == subject_category_id,
self.__model_list[page_no]["layout_dets"],
),
)
)
)
objects = self.__reduct_overlap(
list(
map(
lambda x: x["bbox"],
filter(
lambda x: x["category_id"] == object_category_id,
self.__model_list[page_no]["layout_dets"],
),
)
)
)
subject_object_relation_map = {}
subjects.sort(key=lambda x: x[0] ** 2 + x[1] ** 2) # get the distance !
all_bboxes = []
for v in subjects:
all_bboxes.append({"category_id": subject_category_id, "bbox": v})
for v in objects:
all_bboxes.append({"category_id": object_category_id, "bbox": v})
N = len(all_bboxes)
dis = [[MAX_DIS_OF_POINT] * N for _ in range(N)]
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
):
continue
dis[i][j] = bbox_distance(all_bboxes[i]["bbox"], all_bboxes[j]["bbox"])
dis[j][i] = dis[i][j]
used = set()
for i in range(N):
# 求第 i 个 subject 所关联的 object
if all_bboxes[i]["category_id"] != subject_category_id:
continue
seen = set()
candidates = []
arr = []
for j in range(N):
pos_flag_count = sum(
list(
map(
lambda x: 1 if x else 0,
bbox_relative_pos(
all_bboxes[i]["bbox"], all_bboxes[j]["bbox"]
),
)
)
)
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
):
continue
arr.append((dis[i][j], j))
arr.sort(key=lambda x: x[0])
if len(arr) > 0:
candidates.append(arr[0][1])
seen.add(arr[0][1])
# 已经获取初始种子
for j in set(candidates):
tmp = []
for k in range(i + 1, N):
pos_flag_count = sum(
list(
map(
lambda x: 1 if x else 0,
bbox_relative_pos(
all_bboxes[j]["bbox"], all_bboxes[k]["bbox"]
),
)
)
)
if pos_flag_count > 1:
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
):
continue
is_nearest = True
for l in range(i + 1, N):
if l in (j, k) or l in used or l in seen:
continue
if not float_gt(dis[l][k], dis[j][k]):
is_nearest = False
break
if is_nearest:
tmp.append(k)
seen.add(k)
candidates = tmp
if len(candidates) == 0:
break
# 已经获取到某个 figure 下所有的最靠近的 captions，以及最靠近这些 captions 的 captions 。
# 先扩一下 bbox，
x0s = [all_bboxes[idx]["bbox"][0] for idx in seen] + [
all_bboxes[i]["bbox"][0]
]
y0s = [all_bboxes[idx]["bbox"][1] for idx in seen] + [
all_bboxes[i]["bbox"][1]
]
x1s = [all_bboxes[idx]["bbox"][2] for idx in seen] + [
all_bboxes[i]["bbox"][2]
]
y1s = [all_bboxes[idx]["bbox"][3] for idx in seen] + [
all_bboxes[i]["bbox"][3]
]
ox0, oy0, ox1, oy1 = min(x0s), min(y0s), max(x1s), max(y1s)
ix0, iy0, ix1, iy1 = all_bboxes[i]["bbox"]
# 分成了 4 个截取空间，需要计算落在每个截取空间下 objects 合并后占据的矩形面积
caption_poses = [
[ox0, oy0, ix0, oy1],
[ox0, oy0, ox1, iy0],
[ox0, iy1, ox1, oy1],
[ix1, oy0, ox1, oy1],
]
caption_areas = []
for bbox in caption_poses:
embed_arr = []
for idx in seen:
if _is_in(all_bboxes[idx]["bbox"], bbox):
embed_arr.append(idx)
if len(embed_arr) > 0:
embed_x0 = min([all_bboxes[idx]["bbox"][0] for idx in embed_arr])
embed_y0 = min([all_bboxes[idx]["bbox"][1] for idx in embed_arr])
embed_x1 = max([all_bboxes[idx]["bbox"][2] for idx in embed_arr])
embed_y1 = max([all_bboxes[idx]["bbox"][3] for idx in embed_arr])
caption_areas.append(
int(abs(embed_x1 - embed_x0) * abs(embed_y1 - embed_y0))
)
else:
caption_areas.append(0)
subject_object_relation_map[i] = []
if max(caption_areas) > 0:
max_area_idx = caption_areas.index(max(caption_areas))
caption_bbox = caption_poses[max_area_idx]
for j in seen:
if _is_in(all_bboxes[j]["bbox"], caption_bbox):
used.add(j)
subject_object_relation_map[i].append(j)
for i in sorted(subject_object_relation_map.keys()):
result = {
"subject_body": all_bboxes[i]["bbox"],
"all": all_bboxes[i]["bbox"],
}
if len(subject_object_relation_map[i]) > 0:
x0 = min(
[all_bboxes[j]["bbox"][0] for j in subject_object_relation_map[i]]
)
y0 = min(
[all_bboxes[j]["bbox"][1] for j in subject_object_relation_map[i]]
)
x1 = max(
[all_bboxes[j]["bbox"][2] for j in subject_object_relation_map[i]]
)
y1 = max(
[all_bboxes[j]["bbox"][3] for j in subject_object_relation_map[i]]
)
result["object_body"] = [x0, y0, x1, y1]
result["all"] = [
min(x0, all_bboxes[i]["bbox"][0]),
min(y0, all_bboxes[i]["bbox"][1]),
max(x1, all_bboxes[i]["bbox"][2]),
max(y1, all_bboxes[i]["bbox"][3]),
]
ret.append(result)
total_subject_object_dis = 0
# 计算已经配对的 distance 距离
for i in subject_object_relation_map.keys():
for j in subject_object_relation_map[i]:
total_subject_object_dis += bbox_distance(
all_bboxes[i]["bbox"], all_bboxes[j]["bbox"]
)
# 计算未匹配的 subject 和 object 的距离（非精确版）
with_caption_subject = set(
[
key
for key in subject_object_relation_map.keys()
if len(subject_object_relation_map[i]) > 0
]
)
for i in range(N):
if all_bboxes[i]["category_id"] != object_category_id or i in used:
continue
candidates = []
for j in range(N):
if (
all_bboxes[j]["category_id"] != subject_category_id
or j in with_caption_subject
):
continue
candidates.append((dis[i][j], j))
if len(candidates) > 0:
candidates.sort(key=lambda x: x[0])
total_subject_object_dis += candidates[0][1]
with_caption_subject.add(j)
return ret, total_subject_object_dis
def get_imgs(self, page_no: int): # @许瑞
records, _ = self.__tie_up_category_by_distance(page_no, 3, 4)
return [
{
"bbox": record["all"],
"img_body_bbox": record["subject_body"],
"img_caption_bbox": record.get("object_body", None),
}
for record in records
]
def get_tables(
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)
ret = []
N, M = len(with_captions), len(with_footnotes)
assert N == M
for i in range(N):
record = {
"table_caption_bbox": with_captions[i].get("object_body", None),
"table_body_bbox": with_captions[i]["subject_body"],
"table_footnote_bbox": with_footnotes[i].get("object_body", None),
}
x0 = min(with_captions[i]["all"][0], with_footnotes[i]["all"][0])
y0 = min(with_captions[i]["all"][1], with_footnotes[i]["all"][1])
x1 = max(with_captions[i]["all"][2], with_footnotes[i]["all"][2])
y1 = max(with_captions[i]["all"][3], with_footnotes[i]["all"][3])
record["bbox"] = [x0, y0, x1, y1]
ret.append(record)
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)
return inline_equations, interline_equations, interline_equations_blocks
def get_discarded(self, page_no: int) -> list: # 自研模型，只有坐标
blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.ABANDON.value, page_no)
return blocks
def get_text_blocks(self, page_no: int) -> list: # 自研模型搞的，只有坐标，没有字
blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.PLAIN_TEXT.value, page_no)
return blocks
def get_title_blocks(self, page_no: int) -> list: # 自研模型，只有坐标，没字
blocks = self.__get_blocks_by_type(ModelBlockTypeEnum.TITLE.value, page_no)
return blocks
def get_ocr_text(self, page_no: int) -> list: # paddle 搞的，有字也有坐标
text_spans = []
model_page_info = self.__model_list[page_no]
layout_dets = model_page_info["layout_dets"]
for layout_det in layout_dets:
if layout_det["category_id"] == "15":
span = {
"bbox": layout_det['bbox'],
"content": layout_det["text"],
}
text_spans.append(span)
return text_spans
def get_all_spans(self, page_no: int) -> list:
all_spans = []
model_page_info = self.__model_list[page_no]
layout_dets = model_page_info["layout_dets"]
allow_category_id_list = [3, 5, 13, 14, 15]
"""当成span拼接的"""
# 3: 'image', # 图片
# 4: 'table', # 表格
# 13: 'inline_equation', # 行内公式
# 14: 'interline_equation', # 行间公式
# 15: 'text', # ocr识别文本
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']
}
if category_id == 3:
span["type"] = ContentType.Image
elif category_id == 5:
span["type"] = ContentType.Table
elif category_id == 13:
span["content"] = layout_det["latex"]
span["type"] = ContentType.InlineEquation
elif category_id == 14:
span["content"] = layout_det["latex"]
span["type"] = ContentType.InterlineEquation
elif category_id == 15:
span["content"] = layout_det["text"]
span["type"] = ContentType.Text
all_spans.append(span)
return all_spans
def get_page_size(self, page_no: int): # 获取页面宽高
# 获取当前页的page对象
page = self.__docs[page_no]
# 获取当前页的宽高
page_w = page.rect.width
page_h = page.rect.height
return page_w, page_h
def __get_blocks_by_type(self, types: list, page_no: int, extra_col: list[str] = []) -> list:
blocks = []
for page_dict in self.__model_list:
layout_dets = page_dict.get("layout_dets", [])
page_info = page_dict.get("page_info", {})
page_number = page_info.get("page_no", -1)
if page_no != page_number:
continue
for item in layout_dets:
category_id = item.get("category_id", -1)
bbox = item.get("bbox", None)
if category_id in types:
block = {
"bbox": bbox
}
for col in extra_col:
block[col] = item.get(col, None)
blocks.append(block)
return blocks
if __name__ == "__main__": if __name__ == "__main__":
drw = DiskReaderWriter(r"D:/project/20231108code-clean") drw = DiskReaderWriter(r"D:/project/20231108code-clean")
if 0: if 0:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment