Procházet zdrojové kódy

qa需求定制输出

赵小蒙 před 1 rokem
rodič
revize
ef267e0957
3 změnil soubory, kde provedl 59 přidání a 9 odebrání
  1. 18 8
      demo/ocr_demo.py
  2. 13 0
      magic_pdf/dict2md/ocr_mkcontent.py
  3. 28 1
      magic_pdf/pipeline.py

+ 18 - 8
demo/ocr_demo.py

@@ -4,6 +4,7 @@ import os
 from loguru import logger
 from pathlib import Path
 
+from demo.demo_test import get_json_from_local_or_s3
 from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_mm_markdown_with_para, ocr_mk_nlp_markdown, ocr_mk_mm_markdown, ocr_mk_mm_standard_format
 from magic_pdf.libs.commons import join_path
 from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
@@ -29,14 +30,7 @@ def read_json_file(file_path):
     return data
 
 
-if __name__ == '__main__':
-    #ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.pdf"
-    #ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.json"
-    # ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.pdf"
-    # ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.json"
-    
-    ocr_pdf_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/ocr_1.pdf"
-    ocr_json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/ocr_1.json"
+def ocr_local_parse(ocr_pdf_path, ocr_json_file_path):
     try:
         ocr_pdf_model_info = read_json_file(ocr_json_file_path)
         pth = Path(ocr_json_file_path)
@@ -72,3 +66,19 @@ if __name__ == '__main__':
         # save_markdown(markdown_text, ocr_json_file_path)
     except Exception as e:
         logger.exception(e)
+
+
+def ocr_online_parse(book_name, start_page_id=0, debug_mode=True):
+    json_object = get_json_from_local_or_s3(book_name)
+    logger.info(json_object)
+
+if __name__ == '__main__':
+    #ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.pdf"
+    #ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\s0043-1354(02)00581-x.json"
+    # ocr_pdf_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.pdf"
+    # ocr_json_file_path = r"D:\project\20231108code-clean\ocr\new\双栏\j.1540-627x.2006.00176.x.json"
+    ocr_pdf_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/ocr_1.pdf"
+    ocr_json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/ocr_1.json"
+    ocr_online_parse(book_name="数学新星网/edu_00001236")
+    ocr_local_parse(ocr_pdf_path, ocr_json_file_path)
+    pass

+ 13 - 0
magic_pdf/dict2md/ocr_mkcontent.py

@@ -79,6 +79,19 @@ def ocr_mk_mm_markdown_with_para(pdf_info_dict: dict):
     return '\n\n'.join(markdown)
 
 
+def make_standard_format_with_para(pdf_info_dict: dict):
+    content_list = []
+    for _, page_info in pdf_info_dict.items():
+        paras = page_info.get("para_blocks")
+        if not paras:
+            continue
+        for para in paras:
+            for line in para:
+                content = line_to_standard_format(line)
+                content_list.append(content)
+    return content_list
+
+
 def line_to_standard_format(line):
     line_text = ""
     inline_equation_num = 0

+ 28 - 1
magic_pdf/pipeline.py

@@ -3,7 +3,8 @@ import sys
 import time
 from urllib.parse import quote
 
-from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_nlp_markdown, ocr_mk_mm_markdown, ocr_mk_mm_standard_format
+from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_nlp_markdown, ocr_mk_mm_markdown, ocr_mk_mm_standard_format, \
+    ocr_mk_mm_markdown_with_para
 from magic_pdf.libs.commons import read_file, join_path, parse_bucket_key, formatted_time, s3_image_save_path
 from magic_pdf.libs.drop_reason import DropReason
 from magic_pdf.libs.json_compressor import JsonCompressor
@@ -407,6 +408,32 @@ def ocr_pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
     return jso
 
 
+def ocr_pdf_intermediate_dict_to_markdown_with_para_for_qa(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", file=sys.stderr)
+            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)
+        markdown_content = ocr_mk_mm_markdown_with_para(pdf_intermediate_dict)
+        jso["content_ocr"] = markdown_content
+        logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}", file=sys.stderr)
+        # 把无用的信息清空
+        jso["doc_layout_result"] = ""
+        jso["pdf_intermediate_dict"] = ""
+        jso["pdf_meta"] = ""
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso
+
+
 def ocr_pdf_intermediate_dict_to_standard_format(jso: dict, debug_mode=False) -> dict:
 
     if debug_mode: