magic-pdf 0.6.0__py3-none-any.whl → 0.6.2b1__py3-none-any.whl

magic_pdf/cli/magicpdf.py

@@ -85,10 +85,10 @@ def do_parse(
         f_dump_content_list=True,
         f_make_md_mode=MakeMode.MM_MD,
 ):
+
     orig_model_list = copy.deepcopy(model_list)
 
     local_image_dir, local_md_dir = prepare_env(pdf_file_name, parse_method)
-    logger.info(f"local output dir is {local_md_dir}")
     image_writer, md_writer = DiskReaderWriter(local_image_dir), DiskReaderWriter(local_md_dir)
     image_dir = str(os.path.basename(local_image_dir))
 
@@ -162,6 +162,7 @@ def do_parse(
             path=f"{pdf_file_name}_content_list.json",
             mode=AbsReaderWriter.MODE_TXT,
         )
+    logger.info(f"local output dir is '{local_md_dir}', you can found the result in it.")
 
 
 @click.group()
@@ -179,8 +180,9 @@ def cli():
     help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
     default="auto",
 )
-@click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试")
-@click.option("--model_mode", type=click.STRING, default="full", help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢")
+@click.option("--inside_model", type=click.BOOL, default=True, help="使用内置模型测试")
+@click.option("--model_mode", type=click.STRING, default="full",
+              help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢")
 def json_command(json, method, inside_model, model_mode):
     model_config.__use_inside_model__ = inside_model
     model_config.__model_mode__ = model_mode
@@ -232,8 +234,9 @@ def json_command(json, method, inside_model, model_mode):
     help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
     default="auto",
 )
-@click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试")
-@click.option("--model_mode", type=click.STRING, default="full", help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢")
+@click.option("--inside_model", type=click.BOOL, default=True, help="使用内置模型测试")
+@click.option("--model_mode", type=click.STRING, default="full",
+              help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢")
 def local_json_command(local_json, method, inside_model, model_mode):
     model_config.__use_inside_model__ = inside_model
     model_config.__model_mode__ = model_mode
@@ -277,8 +280,8 @@ def local_json_command(local_json, method, inside_model, model_mode):
 
 @cli.command()
 @click.option(
-    "--pdf", type=click.Path(exists=True), required=True, help="PDF文件的路径"
-)
+    "--pdf", type=click.Path(exists=True), required=True,
+    help='pdf 文件路径, 支持单个文件或文件列表, 文件列表需要以".list"结尾, 一行一个pdf文件路径')
 @click.option("--model", type=click.Path(exists=True), help="模型的路径")
 @click.option(
     "--method",
@@ -286,8 +289,9 @@ def local_json_command(local_json, method, inside_model, model_mode):
     help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
     default="auto",
 )
-@click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试")
-@click.option("--model_mode", type=click.STRING, default="full", help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢")
+@click.option("--inside_model", type=click.BOOL, default=True, help="使用内置模型测试")
+@click.option("--model_mode", type=click.STRING, default="full",
+              help="内置模型选择。lite: 快速解析，精度较低，full: 高精度解析，速度较慢")
 def pdf_command(pdf, model, method, inside_model, model_mode):
     model_config.__use_inside_model__ = inside_model
     model_config.__model_mode__ = model_mode
@@ -296,12 +300,10 @@ def pdf_command(pdf, model, method, inside_model, model_mode):
         disk_rw = DiskReaderWriter(os.path.dirname(path))
         return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
 
-    pdf_data = read_fn(pdf)
-
-    def get_model_json(model_path):
+    def get_model_json(model_path, doc_path):
         # 这里处理pdf和模型相关的逻辑
         if model_path is None:
-            file_name_without_extension, extension = os.path.splitext(pdf)
+            file_name_without_extension, extension = os.path.splitext(doc_path)
             if extension == ".pdf":
                 model_path = file_name_without_extension + ".json"
             else:
@@ -319,15 +321,35 @@ def pdf_command(pdf, model, method, inside_model, model_mode):
 
         return model_json
 
-    jso = json_parse.loads(get_model_json(model))
-    pdf_file_name = Path(pdf).stem
+    def parse_doc(doc_path):
+        try:
+            file_name = str(Path(doc_path).stem)
+            pdf_data = read_fn(doc_path)
+            jso = json_parse.loads(get_model_json(model, doc_path))
 
-    do_parse(
-        pdf_file_name,
-        pdf_data,
-        jso,
-        method,
-    )
+            do_parse(
+                file_name,
+                pdf_data,
+                jso,
+                method,
+            )
+
+        except Exception as e:
+            logger.exception(e)
+
+    if not pdf:
+        logger.error(f"Error: Missing argument '--pdf'.")
+        exit(f"Error: Missing argument '--pdf'.")
+    else:
+        '''适配多个文档的list文件输入'''
+        if pdf.endswith(".list"):
+            with open(pdf, "r") as f:
+                for line in f.readlines():
+                    line = line.strip()
+                    parse_doc(line)
+        else:
+            '''适配单个文档的输入'''
+            parse_doc(pdf)
 
 
 if __name__ == "__main__":

magic_pdf/dict2md/ocr_mkcontent.py

@@ -112,7 +112,7 @@ def ocr_mk_markdown_with_para_core_v2(paras_of_layout, mode, img_buket_path=""):
                         for line in block['lines']:
                             for span in line['spans']:
                                 if span['type'] == ContentType.Image:
-                                    para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
+                                    para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})  \n"
                 for block in para_block['blocks']:  # 2nd.拼image_caption
                     if block['type'] == BlockType.ImageCaption:
                         para_text += merge_para_with_text(block)
@@ -128,7 +128,7 @@ def ocr_mk_markdown_with_para_core_v2(paras_of_layout, mode, img_buket_path=""):
                         for line in block['lines']:
                             for span in line['spans']:
                                 if span['type'] == ContentType.Table:
-                                    para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})\n"
+                                    para_text += f"\n![]({join_path(img_buket_path, span['image_path'])})  \n"
                 for block in para_block['blocks']:  # 3rd.拼table_footnote
                     if block['type'] == BlockType.TableFootnote:
                         para_text += merge_para_with_text(block)
@@ -210,28 +210,32 @@ def para_to_standard_format(para, img_buket_path):
     return para_content
 
 
-def para_to_standard_format_v2(para_block, img_buket_path):
+def para_to_standard_format_v2(para_block, img_buket_path, page_idx):
     para_type = para_block['type']
     if para_type == BlockType.Text:
         para_content = {
             'type': 'text',
             'text': merge_para_with_text(para_block),
+            'page_idx': page_idx
         }
     elif para_type == BlockType.Title:
         para_content = {
             'type': 'text',
             'text': merge_para_with_text(para_block),
-            'text_level': 1
+            'text_level': 1,
+            'page_idx': page_idx
         }
     elif para_type == BlockType.InterlineEquation:
         para_content = {
             'type': 'equation',
             'text': merge_para_with_text(para_block),
-            'text_format': "latex"
+            'text_format': "latex",
+            'page_idx': page_idx
         }
     elif para_type == BlockType.Image:
         para_content = {
             'type': 'image',
+            'page_idx': page_idx
         }
         for block in para_block['blocks']:
             if block['type'] == BlockType.ImageBody:
@@ -241,6 +245,7 @@ def para_to_standard_format_v2(para_block, img_buket_path):
     elif para_type == BlockType.Table:
         para_content = {
             'type': 'table',
+            'page_idx': page_idx
         }
         for block in para_block['blocks']:
             if block['type'] == BlockType.TableBody:
@@ -345,6 +350,7 @@ def union_make(pdf_info_dict: list, make_mode: str, drop_mode: str, img_buket_pa
                 raise Exception(f"drop_mode can not be null")
 
         paras_of_layout = page_info.get("para_blocks")
+        page_idx = page_info.get("page_idx")
         if not paras_of_layout:
             continue
         if make_mode == MakeMode.MM_MD:
@@ -355,7 +361,7 @@ def union_make(pdf_info_dict: list, make_mode: str, drop_mode: str, img_buket_pa
             output_content.extend(page_markdown)
         elif make_mode == MakeMode.STANDARD_FORMAT:
             for para_block in paras_of_layout:
-                para_content = para_to_standard_format_v2(para_block, img_buket_path)
+                para_content = para_to_standard_format_v2(para_block, img_buket_path, page_idx)
                 output_content.append(para_content)
     if make_mode in [MakeMode.MM_MD, MakeMode.NLP_MD]:
         return '\n\n'.join(output_content)

magic_pdf/libs/config_reader.py

@@ -10,16 +10,19 @@ from loguru import logger
 
 from magic_pdf.libs.commons import parse_bucket_key
 
+# 定义配置文件名常量
+CONFIG_FILE_NAME = "magic-pdf.json"
+
 
 def read_config():
     home_dir = os.path.expanduser("~")
 
-    config_file = os.path.join(home_dir, "magic-pdf.json")
+    config_file = os.path.join(home_dir, CONFIG_FILE_NAME)
 
     if not os.path.exists(config_file):
-        raise Exception(f"{config_file} not found")
+        raise FileNotFoundError(f"{config_file} not found")
 
-    with open(config_file, "r") as f:
+    with open(config_file, "r", encoding="utf-8") as f:
         config = json.load(f)
     return config
 
@@ -37,7 +40,7 @@ def get_s3_config(bucket_name: str):
         access_key, secret_key, storage_endpoint = bucket_info[bucket_name]
 
     if access_key is None or secret_key is None or storage_endpoint is None:
-        raise Exception("ak, sk or endpoint not found in magic-pdf.json")
+        raise Exception(f"ak, sk or endpoint not found in {CONFIG_FILE_NAME}")
 
     # logger.info(f"get_s3_config: ak={access_key}, sk={secret_key}, endpoint={storage_endpoint}")
 
@@ -56,17 +59,32 @@ def get_bucket_name(path):
 
 def get_local_dir():
     config = read_config()
-    return config.get("temp-output-dir", "/tmp")
+    local_dir = config.get("temp-output-dir")
+    if local_dir is None:
+        logger.warning(f"'temp-output-dir' not found in {CONFIG_FILE_NAME}, use '/tmp' as default")
+        return "/tmp"
+    else:
+        return local_dir
 
 
 def get_local_models_dir():
     config = read_config()
-    return config.get("models-dir", "/tmp/models")
+    models_dir = config.get("models-dir")
+    if models_dir is None:
+        logger.warning(f"'models-dir' not found in {CONFIG_FILE_NAME}, use '/tmp/models' as default")
+        return "/tmp/models"
+    else:
+        return models_dir
 
 
 def get_device():
     config = read_config()
-    return config.get("device-mode", "cpu")
+    device = config.get("device-mode")
+    if device is None:
+        logger.warning(f"'device-mode' not found in {CONFIG_FILE_NAME}, use 'cpu' as default")
+        return "cpu"
+    else:
+        return device
 
 
 if __name__ == "__main__":

magic_pdf/libs/language.py

@@ -1,8 +1,19 @@
+import os
 import unicodedata
+
+if not os.getenv("FTLANG_CACHE"):
+    current_file_path = os.path.abspath(__file__)
+    current_dir = os.path.dirname(current_file_path)
+    root_dir = os.path.dirname(current_dir)
+    ftlang_cache_dir = os.path.join(root_dir, 'resources', 'fasttext-langdetect')
+    os.environ["FTLANG_CACHE"] = str(ftlang_cache_dir)
+    # print(os.getenv("FTLANG_CACHE"))
+
 from fast_langdetect import detect_language
 
 
 def detect_lang(text: str) -> str:
+
     if len(text) == 0:
         return ""
     try:
@@ -18,6 +29,7 @@ def detect_lang(text: str) -> str:
 
 
 if __name__ == '__main__':
+    print(os.getenv("FTLANG_CACHE"))
     print(detect_lang("This is a test."))
     print(detect_lang("<html>This is a test</html>"))
     print(detect_lang("这个是中文测试。"))

magic_pdf/libs/version.py

@@ -1 +1 @@
-__version__ = "0.6.0"
+__version__ = "0.6.2b1"

magic_pdf/model/__init__.py

@@ -1,2 +1,2 @@
-__use_inside_model__ = False
+__use_inside_model__ = True
 __model_mode__ = "full"

magic_pdf/model/doc_analyze_by_custom_model.py

@@ -48,10 +48,28 @@ def load_images_from_pdf(pdf_bytes: bytes, dpi=200) -> list:
     return images
 
 
-def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False):
+class ModelSingleton:
+    _instance = None
+    _models = {}
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def get_model(self, ocr: bool, show_log: bool):
+        key = (ocr, show_log)
+        if key not in self._models:
+            self._models[key] = custom_model_init(ocr=ocr, show_log=show_log)
+        return self._models[key]
+
+
+def custom_model_init(ocr: bool = False, show_log: bool = False):
     model = None
 
     if model_config.__model_mode__ == "lite":
+        logger.warning("The Lite mode is provided for developers to conduct testing only, and the output quality is "
+                       "not guaranteed to be reliable.")
         model = MODEL.Paddle
     elif model_config.__model_mode__ == "full":
         model = MODEL.PEK
@@ -76,6 +94,14 @@ def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False):
         logger.error("use_inside_model is False, not allow to use inside model")
         exit(1)
 
+    return custom_model
+
+
+def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False):
+
+    model_manager = ModelSingleton()
+    custom_model = model_manager.get_model(ocr, show_log)
+
     images = load_images_from_pdf(pdf_bytes)
 
     model_json = []

magic_pdf/model/magic_model.py

@@ -9,13 +9,14 @@ from magic_pdf.libs.coordinate_transform import get_scale_ratio
 from magic_pdf.libs.ocr_content_type import ContentType
 from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
 from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
-from magic_pdf.libs.math import float_gt
+from magic_pdf.libs.local_math import float_gt
 from magic_pdf.libs.boxbase import (
     _is_in,
     bbox_relative_pos,
     bbox_distance,
     _is_part_overlap,
-    calculate_overlap_area_in_bbox1_area_ratio, calculate_iou,
+    calculate_overlap_area_in_bbox1_area_ratio,
+    calculate_iou,
 )
 from magic_pdf.libs.ModelBlockTypeEnum import ModelBlockTypeEnum
 
@@ -78,9 +79,23 @@ class MagicModel:
                 for layout_det2 in layout_dets:
                     if layout_det1 == layout_det2:
                         continue
-                    if layout_det1["category_id"] in [0,1,2,3,4,5,6,7,8,9] and layout_det2["category_id"] in [0,1,2,3,4,5,6,7,8,9]:
-                        if calculate_iou(layout_det1['bbox'], layout_det2['bbox']) > 0.9:
-                            if layout_det1['score'] < layout_det2['score']:
+                    if layout_det1["category_id"] in [
+                        0,
+                        1,
+                        2,
+                        3,
+                        4,
+                        5,
+                        6,
+                        7,
+                        8,
+                        9,
+                    ] and layout_det2["category_id"] in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]:
+                        if (
+                            calculate_iou(layout_det1["bbox"], layout_det2["bbox"])
+                            > 0.9
+                        ):
+                            if layout_det1["score"] < layout_det2["score"]:
                                 layout_det_need_remove = layout_det1
                             else:
                                 layout_det_need_remove = layout_det2
@@ -97,11 +112,11 @@ class MagicModel:
     def __init__(self, model_list: list, docs: fitz.Document):
         self.__model_list = model_list
         self.__docs = docs
-        '''为所有模型数据添加bbox信息(缩放，poly->bbox)'''
+        """为所有模型数据添加bbox信息(缩放，poly->bbox)"""
         self.__fix_axis()
-        '''删除置信度特别低的模型数据(<0.05),提高质量'''
+        """删除置信度特别低的模型数据(<0.05),提高质量"""
         self.__fix_by_remove_low_confidence()
-        '''删除高iou(>0.9)数据中置信度较低的那个'''
+        """删除高iou(>0.9)数据中置信度较低的那个"""
         self.__fix_by_remove_high_iou_and_low_confidence()
 
     def __reduct_overlap(self, bboxes):
@@ -125,16 +140,6 @@ class MagicModel:
         ret = []
         MAX_DIS_OF_POINT = 10**9 + 7
 
-        def expand_bbox(bbox1, bbox2):
-            x0 = min(bbox1[0], bbox2[0])
-            y0 = min(bbox1[1], bbox2[1])
-            x1 = max(bbox1[2], bbox2[2])
-            y1 = max(bbox1[3], bbox2[3])
-            return [x0, y0, x1, y1]
-
-        def get_bbox_area(bbox):
-            return abs(bbox[2] - bbox[0]) * abs(bbox[3] - bbox[1])
-
         # subject 和 object 的 bbox 会合并成一个大的 bbox （named: merged bbox）。 筛选出所有和 merged bbox 有 overlap 且 overlap 面积大于 object 的面积的 subjects。
         # 再求出筛选出的 subjects 和 object 的最短距离！
         def may_find_other_nearest_bbox(subject_idx, object_idx):
@@ -177,6 +182,13 @@ class MagicModel:
 
             return ret
 
+        def expand_bbbox(idxes):
+            x0s = [all_bboxes[idx]["bbox"][0] for idx in idxes] 
+            y0s = [all_bboxes[idx]["bbox"][1] for idx in idxes] 
+            x1s = [all_bboxes[idx]["bbox"][2] for idx in idxes] 
+            y1s = [all_bboxes[idx]["bbox"][3] for idx in idxes] 
+            return min(x0s), min(y0s), max(x1s), max(y1s)
+
         subjects = self.__reduct_overlap(
             list(
                 map(
@@ -268,7 +280,9 @@ class MagicModel:
                     or dis[i][j] == MAX_DIS_OF_POINT
                 ):
                     continue
-                left, right, _, _ = bbox_relative_pos(all_bboxes[i]["bbox"], all_bboxes[j]["bbox"]) # 由  pos_flag_count 相关逻辑保证本段逻辑准确性
+                left, right, _, _ = bbox_relative_pos(
+                    all_bboxes[i]["bbox"], all_bboxes[j]["bbox"]
+                )  # 由  pos_flag_count 相关逻辑保证本段逻辑准确性
                 if left or right:
                     one_way_dis = all_bboxes[i]["bbox"][2] - all_bboxes[i]["bbox"][0]
                 else:
@@ -322,6 +336,10 @@ class MagicModel:
                             break
 
                     if is_nearest:
+                        nx0, ny0, nx1, ny1 = expand_bbbox(list(seen) + [k])
+                        n_dis = bbox_distance(all_bboxes[i]["bbox"], [nx0, ny0, nx1, ny1])
+                        if float_gt(dis[i][j], n_dis):
+                            continue
                         tmp.append(k)
                         seen.add(k)
 
@@ -331,20 +349,7 @@ class MagicModel:
 
             # 已经获取到某个 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)
+            ox0, oy0, ox1, oy1 = expand_bbbox(list(seen) + [i])
             ix0, iy0, ix1, iy1 = all_bboxes[i]["bbox"]
 
             # 分成了 4 个截取空间，需要计算落在每个截取空间下 objects 合并后占据的矩形面积
@@ -455,8 +460,10 @@ class MagicModel:
                 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)
+    def get_imgs(self, page_no: int):
+        figure_captions, _ = self.__tie_up_category_by_distance(
+            page_no, 3, 4
+        )
         return [
             {
                 "bbox": record["all"],
@@ -464,7 +471,7 @@ class MagicModel:
                 "img_caption_bbox": record.get("object_body", None),
                 "score": record["score"],
             }
-            for record in records
+            for record in figure_captions
         ]
 
     def get_tables(
@@ -535,6 +542,7 @@ class MagicModel:
                 if not any(span == existing_span for existing_span in new_spans):
                     new_spans.append(span)
             return new_spans
+
         all_spans = []
         model_page_info = self.__model_list[page_no]
         layout_dets = model_page_info["layout_dets"]
@@ -548,10 +556,7 @@ class MagicModel:
         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"]
-                }
+                span = {"bbox": layout_det["bbox"], "score": layout_det["score"]}
                 if category_id == 3:
                     span["type"] = ContentType.Image
                 elif category_id == 5:
@@ -604,7 +609,6 @@ class MagicModel:
         return self.__model_list[page_no]
 
 
-
 if __name__ == "__main__":
     drw = DiskReaderWriter(r"D:/project/20231108code-clean")
     if 0: