实现多模态markdown拼装

magic_pdf/pipeline.py

@@ -23,17 +23,11 @@ from loguru import logger
 
 from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
 from magic_pdf.pdf_parse_for_train import parse_pdf_for_train
+from magic_pdf.spark.base import exception_handler, get_data_source
 from magic_pdf.train_utils.convert_to_train_format import convert_to_train_format
 from app.common.s3 import get_s3_config, get_s3_client
 
 
-def exception_handler(jso: dict, e):
-    logger.exception(e)
-    jso["need_drop"] = True
-    jso["drop_reason"] = DropReason.Exception
-    jso["exception"] = f"ERROR: {e}"
-    return jso
-
 
 def get_data_type(jso: dict):
     data_type = jso.get("data_type")
@@ -49,13 +43,6 @@ def get_bookid(jso: dict):
     return book_id
 
 
-def get_data_source(jso: dict):
-    data_source = jso.get("data_source")
-    if data_source is None:
-        data_source = jso.get("file_source")
-    return data_source
-
-
 def meta_scan(jso: dict, doc_layout_check=True) -> dict:
     s3_pdf_path = jso.get("file_location")
     s3_config = get_s3_config(s3_pdf_path)

magic_pdf/pipeline_txt.py

@@ -0,0 +1,68 @@
+"""
+文本型pdf转化为统一清洗格式
+"""
+
+# TODO 移动到spark/目录下
+
+from loguru import logger
+from magic_pdf.dict2md.mkcontent import mk_mm_markdown, mk_universal_format
+from magic_pdf.libs.commons import join_path
+from magic_pdf.libs.json_compressor import JsonCompressor
+from magic_pdf.spark.base import exception_handler, get_data_source
+
+
+def txt_pdf_to_standard_format(jso: dict, debug_mode=False) -> dict:
+    """
+    变成统一的标准格式
+    """
+    if debug_mode:
+        pass
+    else:  # 如果debug没开，则检测是否有needdrop字段
+        if jso.get("need_drop", False):
+            book_name = join_path(get_data_source(jso), jso["file_id"])
+            logger.info(f"book_name is:{book_name} need drop")
+            jso["dropped"] = True
+            return jso
+    try:
+        pdf_intermediate_dict = jso["pdf_intermediate_dict"]
+        # 将 pdf_intermediate_dict 解压
+        pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
+        standard_format = mk_universal_format(pdf_intermediate_dict)
+        jso["content_list"] = standard_format
+        logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']},content_list length is {len(standard_format)}",)
+        # 把无用的信息清空
+        jso["doc_layout_result"] = ""
+        jso["pdf_intermediate_dict"] = ""
+        jso["pdf_meta"] = ""
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso
+
+
+def txt_pdf_to_mm_markdown_format(jso: dict, debug_mode=False) -> dict:
+    """
+    变成多模态的markdown格式
+    """
+    if debug_mode:
+        pass
+    else:  # 如果debug没开，则检测是否有needdrop字段
+        if jso.get("need_drop", False):
+            book_name = join_path(get_data_source(jso), jso["file_id"])
+            logger.info(f"book_name is:{book_name} need drop")
+            jso["dropped"] = True
+            return jso
+    try:
+        pdf_intermediate_dict = jso["pdf_intermediate_dict"]
+        # 将 pdf_intermediate_dict 解压
+        pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
+        standard_format = mk_universal_format(pdf_intermediate_dict)
+        mm_content = mk_mm_markdown(standard_format)
+        jso["content_list"] = mm_content
+        logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']},content_list length is {len(standard_format)}",)
+        # 把无用的信息清空
+        jso["doc_layout_result"] = ""
+        jso["pdf_intermediate_dict"] = ""
+        jso["pdf_meta"] = ""
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso

magic_pdf/spark/base.py

@@ -0,0 +1,21 @@
+
+
+from loguru import logger
+
+from magic_pdf.libs.drop_reason import DropReason
+
+
+def get_data_source(jso: dict):
+    data_source = jso.get("data_source")
+    if data_source is None:
+        data_source = jso.get("file_source")
+    return data_source
+
+
+def exception_handler(jso: dict, e):
+    logger.exception(e)
+    jso["need_drop"] = True
+    jso["drop_reason"] = DropReason.Exception
+    jso["exception"] = f"ERROR: {e}"
+    return jso
+

tests/overall_indicator.py

@@ -46,6 +46,8 @@ def indicator_cal(json_standard,json_test):
     '''批量读取中间生成的json文件'''
     test_inline_equations=[]
     test_interline_equations=[]
+    test_inline_euqations_bboxs=[]
+    test_interline_equations_bboxs=[]
     test_dropped_text_bboxes=[]
     test_dropped_text_tag=[]
     test_dropped_image_bboxes=[]
@@ -58,15 +60,20 @@ def indicator_cal(json_standard,json_test):
         mid_json=pd.DataFrame(i)
         mid_json=mid_json.iloc[:,:-1]
         for j1 in mid_json.loc['inline_equations',:]:
-            page_in=[]
+            page_in_text=[]
+            page_in_bbox=[]
             for k1 in j1:
-                page_in.append(k1['latex_text'])
-            test_inline_equations.append(page_in)
+                page_in_text.append(k1['latex_text'])
+                page_in_bbox.append(k1['bbox'])
+            test_inline_equations.append(page_in_text)
+            test_inline_euqations_bboxs.append(page_in_bbox)
         for j2 in mid_json.loc['interline_equations',:]:
-            page_in=[]
+            page_in_text=[]
+            page_in_bbox=[]
             for k2 in j2:
-                page_in.append(k2['latex_text'])
-            test_interline_equations.append(page_in)
+                page_in_text.append(k2['latex_text'])
+            test_interline_equations.append(page_in_text)
+            test_interline_equations_bboxs.append(page_in_bbox)
 
         for j3 in mid_json.loc['droped_text_block',:]:
             page_in_bbox=[]
@@ -101,6 +108,8 @@ def indicator_cal(json_standard,json_test):
 
     standard_inline_equations=[]
     standard_interline_equations=[]
+    standard_inline_euqations_bboxs=[]
+    standard_interline_equations_bboxs=[]
     standard_dropped_text_bboxes=[]
     standard_dropped_text_tag=[]
     standard_dropped_image_bboxes=[]
@@ -113,15 +122,21 @@ def indicator_cal(json_standard,json_test):
         mid_json=pd.DataFrame(i)
         mid_json=mid_json.iloc[:,:-1]
         for j1 in mid_json.loc['inline_equations',:]:
-            page_in=[]
+            page_in_text=[]
+            page_in_bbox=[]
             for k1 in j1:
-                page_in.append(k1['latex_text'])
-            standard_inline_equations.append(page_in)
+                page_in_text.append(k1['latex_text'])
+                page_in_bbox.append(k1['bbox'])
+            standard_inline_equations.append(page_in_text)
+            standard_inline_euqations_bboxs.append(page_in_bbox)
         for j2 in mid_json.loc['interline_equations',:]:
-            page_in=[]
+            page_in_text=[]
+            page_in_bbox=[]
             for k2 in j2:
-                page_in.append(k2['latex_text'])
-            standard_interline_equations.append(page_in)
+                page_in_text.append(k2['latex_text'])
+                page_in_bbox.append(k2['bbox'])
+            standard_interline_equations.append(page_in_text)
+            standard_interline_equations_bboxs.append(page_in_bbox)
         for j3 in mid_json.loc['droped_text_block',:]:
             page_in_bbox=[]
             page_in_tag=[]
@@ -195,6 +210,9 @@ def indicator_cal(json_standard,json_test):
     inline_equations_edit=np.mean(dis1)
     inline_equations_bleu=np.mean(bleu1)
 
+    '''行内公式bbox匹配相关指标'''
+    inline_equations_bbox_report=bbox_match_indicator(test_inline_euqations_bboxs,standard_inline_euqations_bboxs)
+
 
     '''行间公式编辑距离和bleu'''
     dis2=[]
@@ -217,6 +235,10 @@ def indicator_cal(json_standard,json_test):
     interline_equations_bleu=np.mean(bleu2)
 
 
+    '''行间公式bbox匹配相关指标'''
+    interline_equations_bbox_report=bbox_match_indicator(test_interline_equations_bboxs,standard_interline_equations_bboxs)
+
+
 
 
     '''可以先检查page和bbox数量是否一致'''
@@ -289,87 +311,11 @@ def indicator_cal(json_standard,json_test):
 
     '''dropped_image_block的bbox匹配相关指标'''
     '''有数据格式不一致的问题'''
+    image_block_report=bbox_match_indicator(test_dropped_image_bboxes,standard_dropped_image_bboxes)
     
-    test_image_bbox=[]
-    standard_image_bbox=[]
-    for a,b in zip(test_dropped_image_bboxes,standard_dropped_image_bboxes):
-
-        test_page_bbox=[]
-        standard_page_bbox=[]
-        if len(a)==0 and len(b)==0:
-            pass
-        else:
-            for i in b:
-                if len(i)!=4:
-                    continue
-                else:
-                    judge=0
-                    standard_page_bbox.append(1)
-                    for j in a:
-                        if bbox_offset(i,j):
-                            judge=1
-                            test_page_bbox.append(1)
-                            break
-                    if judge==0:
-                        test_page_bbox.append(0)
-                        
-            diff_num=len(a)+test_page_bbox.count(0)-len(b)
-            if diff_num>0:#有多删的情况出现
-                test_page_bbox.extend([1]*diff_num)
-                standard_page_bbox.extend([0]*diff_num)
-
-          
-            test_image_bbox.extend(test_page_bbox)
-            standard_image_bbox.extend(standard_page_bbox)
-
-    
-    image_block_report = {}
-    image_block_report['accuracy']=metrics.accuracy_score(standard_image_bbox,test_image_bbox)
-    image_block_report['precision']=metrics.precision_score(standard_image_bbox,test_image_bbox)
-    image_block_report['recall']=metrics.recall_score(standard_image_bbox,test_image_bbox)
-    image_block_report['f1_score']=metrics.f1_score(standard_image_bbox,test_image_bbox)
-    
-
-
     
     '''dropped_table_block的bbox匹配相关指标'''
-    test_table_bbox=[]
-    standard_table_bbox=[]
-    for a,b in zip(test_dropped_table_bboxes,standard_dropped_table_bboxes):
-
-        test_page_bbox=[]
-        standard_page_bbox=[]
-        if len(a)==0 and len(b)==0:
-            pass
-        else:
-            for i in b:
-                if len(i)!=4:
-                    continue
-                else:
-                    judge=0
-                    standard_page_bbox.append(1)
-                    for j in a:
-                        if bbox_offset(i,j):
-                            judge=1
-                            test_page_bbox.append(1)
-                            break
-                    if judge==0:
-                        test_page_bbox.append(0)
-                        
-            diff_num=len(a)+test_page_bbox.count(0)-len(b)
-            if diff_num>0:#有多删的情况出现
-                test_page_bbox.extend([1]*diff_num)
-                standard_page_bbox.extend([0]*diff_num)
-
-          
-            test_table_bbox.extend(test_page_bbox)
-            standard_table_bbox.extend(standard_page_bbox)
-
-    table_block_report = {}
-    table_block_report['accuracy']=metrics.accuracy_score(standard_table_bbox,test_table_bbox)
-    table_block_report['precision']=metrics.precision_score(standard_table_bbox,test_table_bbox)
-    table_block_report['recall']=metrics.recall_score(standard_table_bbox,test_table_bbox)
-    table_block_report['f1_score']=metrics.f1_score(standard_table_bbox,test_table_bbox)
+    table_block_report=bbox_match_indicator(test_dropped_table_bboxes,standard_dropped_table_bboxes)
     
    
     '''阅读顺序编辑距离的均值'''
@@ -392,6 +338,8 @@ def indicator_cal(json_standard,json_test):
     output['行间公式平均编辑距离']=[interline_equations_edit]
     output['行内公式平均bleu']=[inline_equations_bleu]
     output['行间公式平均bleu']=[interline_equations_bleu]
+    output['行内公式识别相关指标']=[inline_equations_bbox_report]
+    output['行间公式识别相关指标']=[interline_equations_bbox_report]
     output['阅读顺序平均编辑距离']=[preproc_num_edit]
     output['分段准确率']=[acc_para]
     output['删除的text block的相关指标']=[text_block_report]
@@ -434,6 +382,52 @@ def bbox_offset(b_t,b_s):
         return True
     else:
         return False
+    
+
+'''bbox匹配和对齐函数，输出相关指标'''
+'''输入的是以page为单位的bbox列表'''
+def bbox_match_indicator(test_bbox_list,standard_bbox_list):
+    
+    test_bbox=[]
+    standard_bbox=[]
+    for a,b in zip(test_bbox_list,standard_bbox_list):
+
+        test_page_bbox=[]
+        standard_page_bbox=[]
+        if len(a)==0 and len(b)==0:
+            pass
+        else:
+            for i in b:
+                if len(i)!=4:
+                    continue
+                else:
+                    judge=0
+                    standard_page_bbox.append(1)
+                    for j in a:
+                        if bbox_offset(i,j):
+                            judge=1
+                            test_page_bbox.append(1)
+                            break
+                    if judge==0:
+                        test_page_bbox.append(0)
+                        
+            diff_num=len(a)+test_page_bbox.count(0)-len(b)
+            if diff_num>0:#有多删的情况出现
+                test_page_bbox.extend([1]*diff_num)
+                standard_page_bbox.extend([0]*diff_num)
+
+          
+            test_bbox.extend(test_page_bbox)
+            standard_bbox.extend(standard_page_bbox)
+
+    
+    block_report = {}
+    block_report['accuracy']=metrics.accuracy_score(standard_bbox,test_bbox)
+    block_report['precision']=metrics.precision_score(standard_bbox,test_bbox)
+    block_report['recall']=metrics.recall_score(standard_bbox,test_bbox)
+    block_report['f1_score']=metrics.f1_score(standard_bbox,test_bbox)
+
+    return block_report