目录重构

magic_pdf/pipeline.py

@@ -0,0 +1,294 @@
+# coding=utf8
+import sys
+import time
+from urllib.parse import quote
+
+from magic_pdf.libs.commons import read_file, join_path, parse_bucket_key, formatted_time
+from magic_pdf.libs.drop_reason import DropReason
+from magic_pdf.libs.json_compressor import JsonCompressor
+from magic_pdf.dict2md.mkcontent import mk_nlp_markdown
+from magic_pdf.pdf_parse_by_model import parse_pdf_by_model
+from magic_pdf.filter.pdf_classify_by_type import classify
+from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan
+from loguru import logger
+
+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 meta_scan(jso: dict, doc_layout_check=True) -> dict:
+    s3_pdf_path = jso.get('file_location')
+    s3_config = get_s3_config(s3_pdf_path)
+    if doc_layout_check:
+        if 'doc_layout_result' not in jso:  # 检测json中是存在模型数据，如果没有则需要跳过该pdf
+            jso['need_drop'] = True
+            jso['drop_reason'] = DropReason.MISS_DOC_LAYOUT_RESULT
+            return jso
+    try:
+        data_source = jso.get('data_source')
+        file_id = jso.get('file_id')
+        book_name = data_source + "/" + file_id
+
+        # 首页存在超量drawing问题
+        # special_pdf_list = ['zlib/zlib_21822650']
+        # if book_name in special_pdf_list:
+        #     jso['need_drop'] = True
+        #     jso['drop_reason'] = DropReason.SPECIAL_PDF
+        #     return jso
+
+        start_time = time.time()  # 记录开始时间
+        logger.info(f"book_name is:{book_name},start_time is:{formatted_time(start_time)}", file=sys.stderr)
+        file_content = read_file(s3_pdf_path, s3_config)
+        read_file_time = int(time.time() - start_time)  # 计算执行时间
+
+        start_time = time.time()  # 记录开始时间
+        res = pdf_meta_scan(s3_pdf_path, file_content)
+        if res.get('need_drop', False):  # 如果返回的字典里有need_drop，则提取drop_reason并跳过本次解析
+            jso['need_drop'] = True
+            jso['drop_reason'] = res["drop_reason"]
+        else:  # 正常返回
+            jso['pdf_meta'] = res
+            jso['content'] = ""
+            jso['remark'] = ""
+            jso['data_url'] = ""
+        end_time = time.time()  # 记录结束时间
+        meta_scan_time = int(end_time - start_time)  # 计算执行时间
+        logger.info(f"book_name is:{book_name},end_time is:{formatted_time(end_time)},read_file_time is:{read_file_time},meta_scan_time is:{meta_scan_time}", file=sys.stderr)
+        jso['read_file_time'] = read_file_time
+        jso['meta_scan_time'] = meta_scan_time
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso
+
+
+def classify_by_type(jso: dict, debug_mode=False) -> dict:
+    #检测debug开关
+    if debug_mode:
+        pass
+    else:# 如果debug没开，则检测是否有needdrop字段
+        if jso.get('need_drop', False):
+            return jso
+    # 开始正式逻辑
+    try:
+        pdf_meta = jso.get('pdf_meta')
+        data_source = jso.get('data_source')
+        file_id = jso.get('file_id')
+        book_name = data_source + "/" + file_id
+        total_page = pdf_meta["total_page"]
+        page_width = pdf_meta["page_width_pts"]
+        page_height = pdf_meta["page_height_pts"]
+        img_sz_list = pdf_meta["image_info_per_page"]
+        img_num_list = pdf_meta['imgs_per_page']
+        text_len_list = pdf_meta['text_len_per_page']
+        text_layout_list = pdf_meta['text_layout_per_page']
+        text_language = pdf_meta['text_language']
+        # allow_language = ['zh', 'en']  # 允许的语言,目前只允许简中和英文的
+
+        # if text_language not in allow_language:  # 如果语言不在允许的语言中，则drop
+        #     jso['need_drop'] = True
+        #     jso['drop_reason'] = DropReason.NOT_ALLOW_LANGUAGE
+        #     return jso
+        pdf_path = pdf_meta['pdf_path']
+        is_encrypted = pdf_meta['is_encrypted']
+        is_needs_password = pdf_meta['is_needs_password']
+        if is_encrypted or is_needs_password:  # 加密的，需要密码的，没有页面的，都不处理
+            jso['need_drop'] = True
+            jso['drop_reason'] = DropReason.ENCRYPTED
+        else:
+            start_time = time.time()  # 记录开始时间
+            is_text_pdf, results = classify(pdf_path, total_page, page_width, page_height, img_sz_list, text_len_list, img_num_list, text_layout_list)
+            classify_time = int(time.time() - start_time)  # 计算执行时间
+            if is_text_pdf:
+                pdf_meta['is_text_pdf'] = is_text_pdf
+                jso['pdf_meta'] = pdf_meta
+                jso['classify_time'] = classify_time
+                # print(json.dumps(pdf_meta, ensure_ascii=False))
+
+                allow_language = ['zh', 'en']  # 允许的语言,目前只允许简中和英文的
+                if text_language not in allow_language:  # 如果语言不在允许的语言中，则drop
+                    jso['need_drop'] = True
+                    jso['drop_reason'] = DropReason.NOT_ALLOW_LANGUAGE
+                    return jso
+            else:
+                # 先不drop
+                pdf_meta['is_text_pdf'] = is_text_pdf
+                jso['pdf_meta'] = pdf_meta
+                jso['classify_time'] = classify_time
+                jso['need_drop'] = True
+                jso['drop_reason'] = DropReason.NOT_IS_TEXT_PDF
+                extra_info = {"classify_rules": []}
+                for condition, result in results.items():
+                    if not result:
+                        extra_info["classify_rules"].append(condition)
+                jso['extra_info'] = extra_info
+
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso
+
+
+def save_tables_to_s3(jso: dict, debug_mode=False) -> dict:
+
+    if debug_mode:
+        pass
+    else:# 如果debug没开，则检测是否有needdrop字段
+        if jso.get('need_drop', False):
+            logger.info(f"book_name is:{jso['data_source']}/{jso['file_id']} need drop", file=sys.stderr)
+            jso["dropped"] = True
+            return jso
+    try:
+        data_source = jso.get('data_source')
+        file_id = jso.get('file_id')
+        book_name = data_source + "/" + file_id
+        title = jso.get('title')
+        url_encode_title = quote(title, safe='')
+        if data_source != 'scihub':
+            return jso
+        pdf_intermediate_dict = jso['pdf_intermediate_dict']
+        # 将 pdf_intermediate_dict 解压
+        pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
+        i = 0
+        for page in pdf_intermediate_dict.values():
+            if page.get('tables'):
+                if len(page['tables']) > 0:
+                    j = 0
+                    for table in page['tables']:
+                        if debug_mode:
+                            image_path = join_path("s3://mllm-raw-media/pdf2md_img/", book_name, table['image_path'])
+                        else:
+                            image_path = join_path("s3://mllm-raw-media/pdf2md_img/", table['image_path'])
+
+                        if image_path.endswith('.jpg'):
+                            j += 1
+                            s3_client = get_s3_client(image_path)
+                            bucket_name, bucket_key = parse_bucket_key(image_path)
+                            # 通过s3_client获取图片到内存
+                            image_bytes = s3_client.get_object(Bucket=bucket_name, Key=bucket_key)['Body'].read()
+                            # 保存图片到新的位置
+                            if debug_mode:
+                                new_image_path = join_path("s3://mllm-raw-media/pdf2md_img/table_new/", url_encode_title + "_" + table['image_path'].lstrip('tables/'))
+                            else:
+                                new_image_path = join_path("s3://mllm-raw-media/pdf2md_img/table_new/", url_encode_title + f"_page{i}_{j}.jpg")
+
+                            logger.info(new_image_path, file=sys.stderr)
+                            bucket_name, bucket_key = parse_bucket_key(new_image_path)
+                            s3_client.put_object(Bucket=bucket_name, Key=bucket_key, Body=image_bytes)
+                        else:
+                            continue
+            i += 1
+
+        # 把无用的信息清空
+        jso["doc_layout_result"] = ""
+        jso["pdf_intermediate_dict"] = ""
+        jso["pdf_meta"] = ""
+    except Exception as e:
+        jso = exception_handler(jso, e)
+    return jso
+
+
+def drop_needdrop_pdf(jso: dict) -> dict:
+    if jso.get('need_drop', False):
+        logger.info(f"book_name is:{jso['data_source']}/{jso['file_id']} need drop", file=sys.stderr)
+        jso["dropped"] = True
+    return jso
+
+
+def pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
+
+    if debug_mode:
+        pass
+    else:# 如果debug没开，则检测是否有needdrop字段
+        if jso.get('need_drop', False):
+            book_name = join_path(jso['data_source'], 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 = mk_nlp_markdown(pdf_intermediate_dict)
+        jso["content"] = markdown_content
+        logger.info(f"book_name is:{jso['data_source']}/{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 parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
+    #检测debug开关
+    if debug_mode:
+        pass
+    else:# 如果debug没开，则检测是否有needdrop字段
+        if jso.get('need_drop', False):
+            return jso
+    # 开始正式逻辑
+    s3_pdf_path = jso.get('file_location')
+    s3_config = get_s3_config(s3_pdf_path)
+    model_output_json_list = jso.get('doc_layout_result')
+    data_source = jso.get('data_source')
+    file_id = jso.get('file_id')
+    book_name = data_source + "/" + file_id
+
+    # 1.23.22已修复
+    # if debug_mode:
+    #     pass
+    # else:
+    #     if book_name == "zlib/zlib_21929367":
+    #         jso['need_drop'] = True
+    #         jso['drop_reason'] = DropReason.SPECIAL_PDF
+    #         return jso
+
+    junk_img_bojids = jso['pdf_meta']['junk_img_bojids']
+    # total_page = jso['pdf_meta']['total_page']
+
+    # 增加检测 max_svgs 数量的检测逻辑，如果 max_svgs 超过3000则drop
+    svgs_per_page_list = jso['pdf_meta']['svgs_per_page']
+    max_svgs = max(svgs_per_page_list)
+    if max_svgs > 3000:
+        jso['need_drop'] = True
+        jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
+    # elif total_page > 1000:
+    #     jso['need_drop'] = True
+    #     jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_TOTAL_PAGES
+    else:
+        try:
+            save_path = "s3://mllm-raw-media/pdf2md_img/"
+            image_s3_config = get_s3_config(save_path)
+            start_time = time.time()  # 记录开始时间
+            # 先打印一下book_name和解析开始的时间
+            logger.info(f"book_name is:{book_name},start_time is:{formatted_time(start_time)}", file=sys.stderr)
+            pdf_info_dict = parse_pdf_by_model(s3_pdf_path, s3_config, model_output_json_list, save_path,
+                                                  book_name, pdf_model_profile=None,
+                                                  image_s3_config=image_s3_config,
+                                                  start_page_id=start_page_id, junk_img_bojids=junk_img_bojids,
+                                                  debug_mode=debug_mode)
+            if pdf_info_dict.get('need_drop', False):  # 如果返回的字典里有need_drop，则提取drop_reason并跳过本次解析
+                jso['need_drop'] = True
+                jso['drop_reason'] = pdf_info_dict["drop_reason"]
+            else:  # 正常返回，将 pdf_info_dict 压缩并存储
+                pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
+                jso['pdf_intermediate_dict'] = pdf_info_dict
+            end_time = time.time()  # 记录完成时间
+            parse_time = int(end_time - start_time)  # 计算执行时间
+            # 解析完成后打印一下book_name和耗时
+            logger.info(f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}", file=sys.stderr)
+            jso['parse_time'] = parse_time
+        except Exception as e:
+            jso = exception_handler(jso, e)
+    return jso
+
+
+if __name__ == "__main__":
+    pass