373 lines
15 KiB
Python
373 lines
15 KiB
Python
from mineru.utils.boxbase import bbox_relative_pos, calculate_iou, bbox_distance, get_minbox_if_overlap_by_ratio
|
||
from mineru.utils.enum_class import CategoryId, ContentType
|
||
from mineru.utils.magic_model_utils import tie_up_category_by_distance_v3, reduct_overlap
|
||
|
||
|
||
class MagicModel:
|
||
"""每个函数没有得到元素的时候返回空list."""
|
||
def __init__(self, page_model_info: dict, scale: float):
|
||
self.__page_model_info = page_model_info
|
||
self.__scale = scale
|
||
"""为所有模型数据添加bbox信息(缩放，poly->bbox)"""
|
||
self.__fix_axis()
|
||
"""删除置信度特别低的模型数据(<0.05),提高质量"""
|
||
self.__fix_by_remove_low_confidence()
|
||
"""删除高iou(>0.9)数据中置信度较低的那个"""
|
||
self.__fix_by_remove_high_iou_and_low_confidence()
|
||
"""将部分tbale_footnote修正为image_footnote"""
|
||
self.__fix_footnote()
|
||
"""处理重叠的image_body和table_body"""
|
||
self.__fix_by_remove_overlap_image_table_body()
|
||
|
||
def __fix_by_remove_overlap_image_table_body(self):
|
||
need_remove_list = []
|
||
layout_dets = self.__page_model_info['layout_dets']
|
||
image_blocks = list(filter(
|
||
lambda x: x['category_id'] == CategoryId.ImageBody, layout_dets
|
||
))
|
||
table_blocks = list(filter(
|
||
lambda x: x['category_id'] == CategoryId.TableBody, layout_dets
|
||
))
|
||
|
||
def add_need_remove_block(blocks):
|
||
for i in range(len(blocks)):
|
||
for j in range(i + 1, len(blocks)):
|
||
block1 = blocks[i]
|
||
block2 = blocks[j]
|
||
overlap_box = get_minbox_if_overlap_by_ratio(
|
||
block1['bbox'], block2['bbox'], 0.8
|
||
)
|
||
if overlap_box is not None:
|
||
# 判断哪个区块的面积更小，移除较小的区块
|
||
area1 = (block1['bbox'][2] - block1['bbox'][0]) * (block1['bbox'][3] - block1['bbox'][1])
|
||
area2 = (block2['bbox'][2] - block2['bbox'][0]) * (block2['bbox'][3] - block2['bbox'][1])
|
||
|
||
if area1 <= area2:
|
||
block_to_remove = block1
|
||
large_block = block2
|
||
else:
|
||
block_to_remove = block2
|
||
large_block = block1
|
||
|
||
if block_to_remove not in need_remove_list:
|
||
# 扩展大区块的边界框
|
||
x1, y1, x2, y2 = large_block['bbox']
|
||
sx1, sy1, sx2, sy2 = block_to_remove['bbox']
|
||
x1 = min(x1, sx1)
|
||
y1 = min(y1, sy1)
|
||
x2 = max(x2, sx2)
|
||
y2 = max(y2, sy2)
|
||
large_block['bbox'] = [x1, y1, x2, y2]
|
||
need_remove_list.append(block_to_remove)
|
||
|
||
# 处理图像-图像重叠
|
||
add_need_remove_block(image_blocks)
|
||
# 处理表格-表格重叠
|
||
add_need_remove_block(table_blocks)
|
||
|
||
# 从布局中移除标记的区块
|
||
for need_remove in need_remove_list:
|
||
if need_remove in layout_dets:
|
||
layout_dets.remove(need_remove)
|
||
|
||
|
||
def __fix_axis(self):
|
||
need_remove_list = []
|
||
layout_dets = self.__page_model_info['layout_dets']
|
||
for layout_det in layout_dets:
|
||
x0, y0, _, _, x1, y1, _, _ = layout_det['poly']
|
||
bbox = [
|
||
int(x0 / self.__scale),
|
||
int(y0 / self.__scale),
|
||
int(x1 / self.__scale),
|
||
int(y1 / self.__scale),
|
||
]
|
||
layout_det['bbox'] = bbox
|
||
# 删除高度或者宽度小于等于0的spans
|
||
if bbox[2] - bbox[0] <= 0 or bbox[3] - bbox[1] <= 0:
|
||
need_remove_list.append(layout_det)
|
||
for need_remove in need_remove_list:
|
||
layout_dets.remove(need_remove)
|
||
|
||
def __fix_by_remove_low_confidence(self):
|
||
need_remove_list = []
|
||
layout_dets = self.__page_model_info['layout_dets']
|
||
for layout_det in layout_dets:
|
||
if layout_det['score'] <= 0.05:
|
||
need_remove_list.append(layout_det)
|
||
else:
|
||
continue
|
||
for need_remove in need_remove_list:
|
||
layout_dets.remove(need_remove)
|
||
|
||
def __fix_by_remove_high_iou_and_low_confidence(self):
|
||
need_remove_list = []
|
||
layout_dets = list(filter(
|
||
lambda x: x['category_id'] in [
|
||
CategoryId.Title,
|
||
CategoryId.Text,
|
||
CategoryId.ImageBody,
|
||
CategoryId.ImageCaption,
|
||
CategoryId.TableBody,
|
||
CategoryId.TableCaption,
|
||
CategoryId.TableFootnote,
|
||
CategoryId.InterlineEquation_Layout,
|
||
CategoryId.InterlineEquationNumber_Layout,
|
||
], self.__page_model_info['layout_dets']
|
||
)
|
||
)
|
||
for i in range(len(layout_dets)):
|
||
for j in range(i + 1, len(layout_dets)):
|
||
layout_det1 = layout_dets[i]
|
||
layout_det2 = layout_dets[j]
|
||
|
||
if calculate_iou(layout_det1['bbox'], layout_det2['bbox']) > 0.9:
|
||
|
||
layout_det_need_remove = layout_det1 if layout_det1['score'] < layout_det2['score'] else layout_det2
|
||
|
||
if layout_det_need_remove not in need_remove_list:
|
||
need_remove_list.append(layout_det_need_remove)
|
||
|
||
for need_remove in need_remove_list:
|
||
self.__page_model_info['layout_dets'].remove(need_remove)
|
||
|
||
def __fix_footnote(self):
|
||
footnotes = []
|
||
figures = []
|
||
tables = []
|
||
|
||
for obj in self.__page_model_info['layout_dets']:
|
||
if obj['category_id'] == CategoryId.TableFootnote:
|
||
footnotes.append(obj)
|
||
elif obj['category_id'] == CategoryId.ImageBody:
|
||
figures.append(obj)
|
||
elif obj['category_id'] == CategoryId.TableBody:
|
||
tables.append(obj)
|
||
if len(footnotes) * len(figures) == 0:
|
||
continue
|
||
dis_figure_footnote = {}
|
||
dis_table_footnote = {}
|
||
|
||
for i in range(len(footnotes)):
|
||
for j in range(len(figures)):
|
||
pos_flag_count = sum(
|
||
list(
|
||
map(
|
||
lambda x: 1 if x else 0,
|
||
bbox_relative_pos(
|
||
footnotes[i]['bbox'], figures[j]['bbox']
|
||
),
|
||
)
|
||
)
|
||
)
|
||
if pos_flag_count > 1:
|
||
continue
|
||
dis_figure_footnote[i] = min(
|
||
self._bbox_distance(figures[j]['bbox'], footnotes[i]['bbox']),
|
||
dis_figure_footnote.get(i, float('inf')),
|
||
)
|
||
for i in range(len(footnotes)):
|
||
for j in range(len(tables)):
|
||
pos_flag_count = sum(
|
||
list(
|
||
map(
|
||
lambda x: 1 if x else 0,
|
||
bbox_relative_pos(
|
||
footnotes[i]['bbox'], tables[j]['bbox']
|
||
),
|
||
)
|
||
)
|
||
)
|
||
if pos_flag_count > 1:
|
||
continue
|
||
|
||
dis_table_footnote[i] = min(
|
||
self._bbox_distance(tables[j]['bbox'], footnotes[i]['bbox']),
|
||
dis_table_footnote.get(i, float('inf')),
|
||
)
|
||
for i in range(len(footnotes)):
|
||
if i not in dis_figure_footnote:
|
||
continue
|
||
if dis_table_footnote.get(i, float('inf')) > dis_figure_footnote[i]:
|
||
footnotes[i]['category_id'] = CategoryId.ImageFootnote
|
||
|
||
def _bbox_distance(self, bbox1, bbox2):
|
||
left, right, bottom, top = bbox_relative_pos(bbox1, bbox2)
|
||
flags = [left, right, bottom, top]
|
||
count = sum([1 if v else 0 for v in flags])
|
||
if count > 1:
|
||
return float('inf')
|
||
if left or right:
|
||
l1 = bbox1[3] - bbox1[1]
|
||
l2 = bbox2[3] - bbox2[1]
|
||
else:
|
||
l1 = bbox1[2] - bbox1[0]
|
||
l2 = bbox2[2] - bbox2[0]
|
||
|
||
if l2 > l1 and (l2 - l1) / l1 > 0.3:
|
||
return float('inf')
|
||
|
||
return bbox_distance(bbox1, bbox2)
|
||
|
||
def __tie_up_category_by_distance_v3(self, subject_category_id, object_category_id):
|
||
# 定义获取主体和客体对象的函数
|
||
def get_subjects():
|
||
return reduct_overlap(
|
||
list(
|
||
map(
|
||
lambda x: {'bbox': x['bbox'], 'score': x['score']},
|
||
filter(
|
||
lambda x: x['category_id'] == subject_category_id,
|
||
self.__page_model_info['layout_dets'],
|
||
),
|
||
)
|
||
)
|
||
)
|
||
|
||
def get_objects():
|
||
return reduct_overlap(
|
||
list(
|
||
map(
|
||
lambda x: {'bbox': x['bbox'], 'score': x['score']},
|
||
filter(
|
||
lambda x: x['category_id'] == object_category_id,
|
||
self.__page_model_info['layout_dets'],
|
||
),
|
||
)
|
||
)
|
||
)
|
||
|
||
# 调用通用方法
|
||
return tie_up_category_by_distance_v3(
|
||
get_subjects,
|
||
get_objects
|
||
)
|
||
|
||
def get_imgs(self):
|
||
with_captions = self.__tie_up_category_by_distance_v3(
|
||
CategoryId.ImageBody, CategoryId.ImageCaption
|
||
)
|
||
with_footnotes = self.__tie_up_category_by_distance_v3(
|
||
CategoryId.ImageBody, CategoryId.ImageFootnote
|
||
)
|
||
ret = []
|
||
for v in with_captions:
|
||
record = {
|
||
'image_body': v['sub_bbox'],
|
||
'image_caption_list': v['obj_bboxes'],
|
||
}
|
||
filter_idx = v['sub_idx']
|
||
d = next(filter(lambda x: x['sub_idx'] == filter_idx, with_footnotes))
|
||
record['image_footnote_list'] = d['obj_bboxes']
|
||
ret.append(record)
|
||
return ret
|
||
|
||
def get_tables(self) -> list:
|
||
with_captions = self.__tie_up_category_by_distance_v3(
|
||
CategoryId.TableBody, CategoryId.TableCaption
|
||
)
|
||
with_footnotes = self.__tie_up_category_by_distance_v3(
|
||
CategoryId.TableBody, CategoryId.TableFootnote
|
||
)
|
||
ret = []
|
||
for v in with_captions:
|
||
record = {
|
||
'table_body': v['sub_bbox'],
|
||
'table_caption_list': v['obj_bboxes'],
|
||
}
|
||
filter_idx = v['sub_idx']
|
||
d = next(filter(lambda x: x['sub_idx'] == filter_idx, with_footnotes))
|
||
record['table_footnote_list'] = d['obj_bboxes']
|
||
ret.append(record)
|
||
return ret
|
||
|
||
def get_equations(self) -> tuple[list, list, list]: # 有坐标，也有字
|
||
inline_equations = self.__get_blocks_by_type(
|
||
CategoryId.InlineEquation, ['latex']
|
||
)
|
||
interline_equations = self.__get_blocks_by_type(
|
||
CategoryId.InterlineEquation_YOLO, ['latex']
|
||
)
|
||
interline_equations_blocks = self.__get_blocks_by_type(
|
||
CategoryId.InterlineEquation_Layout
|
||
)
|
||
return inline_equations, interline_equations, interline_equations_blocks
|
||
|
||
def get_discarded(self) -> list: # 自研模型，只有坐标
|
||
blocks = self.__get_blocks_by_type(CategoryId.Abandon)
|
||
return blocks
|
||
|
||
def get_text_blocks(self) -> list: # 自研模型搞的，只有坐标，没有字
|
||
blocks = self.__get_blocks_by_type(CategoryId.Text)
|
||
return blocks
|
||
|
||
def get_title_blocks(self) -> list: # 自研模型，只有坐标，没字
|
||
blocks = self.__get_blocks_by_type(CategoryId.Title)
|
||
return blocks
|
||
|
||
def get_all_spans(self) -> list:
|
||
|
||
def remove_duplicate_spans(spans):
|
||
new_spans = []
|
||
for span in spans:
|
||
if not any(span == existing_span for existing_span in new_spans):
|
||
new_spans.append(span)
|
||
return new_spans
|
||
|
||
all_spans = []
|
||
layout_dets = self.__page_model_info['layout_dets']
|
||
allow_category_id_list = [
|
||
CategoryId.ImageBody,
|
||
CategoryId.TableBody,
|
||
CategoryId.InlineEquation,
|
||
CategoryId.InterlineEquation_YOLO,
|
||
CategoryId.OcrText,
|
||
]
|
||
"""当成span拼接的"""
|
||
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'], 'score': layout_det['score']}
|
||
if category_id == CategoryId.ImageBody:
|
||
span['type'] = ContentType.IMAGE
|
||
elif category_id == CategoryId.TableBody:
|
||
# 获取table模型结果
|
||
latex = layout_det.get('latex', None)
|
||
html = layout_det.get('html', None)
|
||
if latex:
|
||
span['latex'] = latex
|
||
elif html:
|
||
span['html'] = html
|
||
span['type'] = ContentType.TABLE
|
||
elif category_id == CategoryId.InlineEquation:
|
||
span['content'] = layout_det['latex']
|
||
span['type'] = ContentType.INLINE_EQUATION
|
||
elif category_id == CategoryId.InterlineEquation_YOLO:
|
||
span['content'] = layout_det['latex']
|
||
span['type'] = ContentType.INTERLINE_EQUATION
|
||
elif category_id == CategoryId.OcrText:
|
||
span['content'] = layout_det['text']
|
||
span['type'] = ContentType.TEXT
|
||
all_spans.append(span)
|
||
return remove_duplicate_spans(all_spans)
|
||
|
||
def __get_blocks_by_type(
|
||
self, category_type: int, extra_col=None
|
||
) -> list:
|
||
if extra_col is None:
|
||
extra_col = []
|
||
blocks = []
|
||
layout_dets = self.__page_model_info.get('layout_dets', [])
|
||
for item in layout_dets:
|
||
category_id = item.get('category_id', -1)
|
||
bbox = item.get('bbox', None)
|
||
|
||
if category_id == category_type:
|
||
block = {
|
||
'bbox': bbox,
|
||
'score': item.get('score'),
|
||
}
|
||
for col in extra_col:
|
||
block[col] = item.get(col, None)
|
||
blocks.append(block)
|
||
return blocks |