magic-pdf 0.5.7__py3-none-any.whl → 0.5.8__py3-none-any.whl

magic_pdf/cli/magicpdf.py

@@ -50,9 +50,7 @@ parse_pdf_methods = click.Choice(["ocr", "txt", "auto"])
 
 
 def prepare_env(pdf_file_name, method):
-    local_parent_dir = os.path.join(
-        get_local_dir(), "magic-pdf", pdf_file_name, method
-    )
+    local_parent_dir = os.path.join(get_local_dir(), "magic-pdf", pdf_file_name, method)
 
     local_image_dir = os.path.join(str(local_parent_dir), "images")
     local_md_dir = local_parent_dir
@@ -62,7 +60,7 @@ def prepare_env(pdf_file_name, method):
 
 
 def write_to_csv(csv_file_path, csv_data):
-    with open(csv_file_path, mode='a', newline='', encoding='utf-8') as csvfile:
+    with open(csv_file_path, mode="a", newline="", encoding="utf-8") as csvfile:
         # 创建csv writer对象
         csv_writer = csv.writer(csvfile)
         # 写入数据
@@ -70,65 +68,92 @@ def write_to_csv(csv_file_path, csv_data):
     print(f"数据已成功追加到 '{csv_file_path}'")
 
 
-def _do_parse(pdf_file_name, pdf_bytes, model_list, parse_method, image_writer, md_writer, image_dir, local_md_dir):
+def do_parse(
+    pdf_file_name,
+    pdf_bytes,
+    model_list,
+    parse_method,
+    f_draw_span_bbox=True,
+    f_draw_layout_bbox=True,
+    f_dump_md=True,
+    f_dump_middle_json=True,
+    f_dump_model_json=True,
+    f_dump_orig_pdf=True,
+    f_dump_content_list=True,
+):
+
+    local_image_dir, local_md_dir = prepare_env(pdf_file_name, parse_method)
+    image_writer, md_writer = DiskReaderWriter(local_image_dir), DiskReaderWriter(local_md_dir)
+    image_dir = (os.path.basename(local_image_dir),)
+
     if parse_method == "auto":
-        jso_useful_key = {
-            "_pdf_type": "",
-            "model_list": model_list
-        }
+        jso_useful_key = {"_pdf_type": "", "model_list": model_list}
         pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True)
     elif parse_method == "txt":
         pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True)
     elif parse_method == "ocr":
         pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True)
     else:
-        print("unknow parse method")
+        print("unknown parse method")
         sys.exit(1)
 
     pipe.pipe_classify()
 
-    '''如果没有传入有效的模型数据，则使用内置paddle解析'''
+    """如果没有传入有效的模型数据，则使用内置paddle解析"""
     if len(model_list) == 0:
         pipe.pipe_analyze()
 
     pipe.pipe_parse()
-    pdf_info = pipe.pdf_mid_data['pdf_info']
-    draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir)
-    draw_span_bbox(pdf_info, pdf_bytes, local_md_dir)
+    pdf_info = pipe.pdf_mid_data["pdf_info"]
+    if f_draw_layout_bbox:
+        draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir)
+    if f_draw_span_bbox:
+        draw_span_bbox(pdf_info, pdf_bytes, local_md_dir)
 
     # write_to_csv(r"D:\project\20231108code-clean\linshixuqiu\pdf_dev\新模型\新建文件夹\luanma.csv",
     #              [pdf_file_name, pipe.pdf_mid_data['not_common_character_rate'], pipe.pdf_mid_data['not_printable_rate']])
 
-    md_content = pipe.pipe_mk_markdown(image_dir, drop_mode=DropMode.NONE)
-    '''写markdown'''
-    md_writer.write(
-        content=md_content, path=f"{pdf_file_name}.md", mode=AbsReaderWriter.MODE_TXT
-    )
-    '''写middle_json'''
-    md_writer.write(
-        content=json_parse.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4),
-        path=f"{pdf_file_name}_middle.json",
-        mode=AbsReaderWriter.MODE_TXT,
-    )
-    '''写model_json'''
-    md_writer.write(
-        content=json_parse.dumps(pipe.model_list, ensure_ascii=False, indent=4),
-        path=f"{pdf_file_name}_model.json",
-        mode=AbsReaderWriter.MODE_TXT,
-    )
-    '''写源pdf'''
-    md_writer.write(
-        content=pdf_bytes,
-        path=f"{pdf_file_name}_origin.pdf",
-        mode=AbsReaderWriter.MODE_BIN,
-    )
-    content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE)
-    '''写content_list'''
-    md_writer.write(
-        content=json_parse.dumps(content_list, ensure_ascii=False, indent=4),
-        path=f"{pdf_file_name}_content_list.json",
-        mode=AbsReaderWriter.MODE_TXT
-    )
+    md_content = pipe.pipe_mk_markdown(str(image_dir), drop_mode=DropMode.NONE)
+    if f_dump_md:
+        """写markdown"""
+        md_writer.write(
+            content=md_content,
+            path=f"{pdf_file_name}.md",
+            mode=AbsReaderWriter.MODE_TXT,
+        )
+
+    if f_dump_middle_json:
+        """写middle_json"""
+        md_writer.write(
+            content=json_parse.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4),
+            path=f"{pdf_file_name}_middle.json",
+            mode=AbsReaderWriter.MODE_TXT,
+        )
+
+    if f_dump_model_json:
+        """写model_json"""
+        md_writer.write(
+            content=json_parse.dumps(pipe.model_list, ensure_ascii=False, indent=4),
+            path=f"{pdf_file_name}_model.json",
+            mode=AbsReaderWriter.MODE_TXT,
+        )
+
+    if f_dump_orig_pdf:
+        """写源pdf"""
+        md_writer.write(
+            content=pdf_bytes,
+            path=f"{pdf_file_name}_origin.pdf",
+            mode=AbsReaderWriter.MODE_BIN,
+        )
+
+    content_list = pipe.pipe_mk_uni_format(str(image_dir), drop_mode=DropMode.NONE)
+    if f_dump_content_list:
+        """写content_list"""
+        md_writer.write(
+            content=json_parse.dumps(content_list, ensure_ascii=False, indent=4),
+            path=f"{pdf_file_name}_content_list.json",
+            mode=AbsReaderWriter.MODE_TXT,
+        )
 
 
 @click.group()
@@ -177,21 +202,12 @@ def json_command(json, method):
         s3_file_path = jso.get("path")
     pdf_file_name = Path(s3_file_path).stem
     pdf_data = read_s3_path(s3_file_path)
-    local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)
-
-    local_image_rw, local_md_rw = DiskReaderWriter(local_image_dir), DiskReaderWriter(
-        local_md_dir
-    )
 
-    _do_parse(
+    do_parse(
         pdf_file_name,
         pdf_data,
         jso["doc_layout_result"],
         method,
-        local_image_rw,
-        local_md_rw,
-        os.path.basename(local_image_dir),
-        local_md_dir
     )
 
 
@@ -233,21 +249,11 @@ def local_json_command(local_json, method):
                 s3_file_path = jso.get("path")
             pdf_file_name = Path(s3_file_path).stem
             pdf_data = read_s3_path(s3_file_path)
-            local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)
-
-            local_image_rw, local_md_rw = DiskReaderWriter(local_image_dir), DiskReaderWriter(
-                local_md_dir
-            )
-
-            _do_parse(
+            do_parse(
                 pdf_file_name,
                 pdf_data,
                 jso["doc_layout_result"],
                 method,
-                local_image_rw,
-                local_md_rw,
-                os.path.basename(local_image_dir),
-                local_md_dir
             )
 
 
@@ -274,7 +280,9 @@ def pdf_command(pdf, model, method):
         if model_path is None:
             model_path = pdf.replace(".pdf", ".json")
             if not os.path.exists(model_path):
-                logger.warning(f"not found json {model_path} existed, use paddle analyze")
+                logger.warning(
+                    f"not found json {model_path} existed, use paddle analyze"
+                )
                 # 本地无模型数据则调用内置paddle分析，先传空list，在内部识别到空list再调用paddle
                 model_json = "[]"
             else:
@@ -286,19 +294,12 @@ def pdf_command(pdf, model, method):
 
     jso = json_parse.loads(get_model_json(model))
     pdf_file_name = Path(pdf).stem
-    local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)
-    local_image_rw, local_md_rw = DiskReaderWriter(local_image_dir), DiskReaderWriter(
-        local_md_dir
-    )
-    _do_parse(
+
+    do_parse(
         pdf_file_name,
         pdf_data,
         jso,
         method,
-        local_image_rw,
-        local_md_rw,
-        os.path.basename(local_image_dir),
-        local_md_dir
     )
 
 

magic_pdf/libs/version.py

@@ -1 +1 @@
-__version__ = "0.5.7"
+__version__ = "0.5.8"

{magic_pdf-0.5.7.dist-info → magic_pdf-0.5.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: magic-pdf
-Version: 0.5.7
+Version: 0.5.8
 Requires-Python: >=3.9
 License-File: LICENSE.md
 Requires-Dist: boto3 >=1.28.43

{magic_pdf-0.5.7.dist-info → magic_pdf-0.5.8.dist-info}/RECORD

@@ -7,7 +7,7 @@ magic_pdf/pdf_parse_for_train.py,sha256=Oby61DMjJ716Jj_ri7lwXfv2Chus0pbBR2RPXrmB
 magic_pdf/pdf_parse_union_core.py,sha256=a67iQuEfuslAEF-wQplGZKXUuz5mT3HiCyvuR52E6Gw,10584
 magic_pdf/user_api.py,sha256=R4onPBXlA8GARwlqBh5wmxUtTxwQb-PUeFK7eTAWzoU,4971
 magic_pdf/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-magic_pdf/cli/magicpdf.py,sha256=IoyuWsnJp5lLDS4G9brtCqNdIWKb57Ini4uftkCl2Mg,11357
+magic_pdf/cli/magicpdf.py,sha256=aMmY_J83_8IdkkxyZUHg8WzIbCA_oW4cyjlUUZE0Wvc,11117
 magic_pdf/dict2md/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 magic_pdf/dict2md/mkcontent.py,sha256=rWUY-2opd0jeowEUEVOV_uWcKum1Q7ng4nOoT6-ka_s,17459
 magic_pdf/dict2md/ocr_mkcontent.py,sha256=RyxebPtvFfNce_HCa-_YGxwFx_srzL-BfMKc85V9JG0,15442
@@ -46,7 +46,7 @@ magic_pdf/libs/pdf_check.py,sha256=MAe8wzwT0qvPf_I72wEZG7k1g4haNHS7oUtLqkB5rlE,2
 magic_pdf/libs/pdf_image_tools.py,sha256=CAd01giTKr_UJz1_QtDOARG9G9z69GFpzRZwcWSfLtE,1282
 magic_pdf/libs/safe_filename.py,sha256=ckwcM_eqoysTb5id8czp-tXq2G9da0-l3pshZDCHQtE,236
 magic_pdf/libs/textbase.py,sha256=SC1Frhz3Fb7V7n2SFRBsl7Bmg0JZdlvZskq0lfW1vIk,732
-magic_pdf/libs/version.py,sha256=KiyyYbyEe0O858kmiWcg1RdmqGUYtk_JqRmc3_Ev2Q8,22
+magic_pdf/libs/version.py,sha256=bDuZ37zImJZsQ3a4pW87q4kg-zsIBrUFAv1aumIf_7k,22
 magic_pdf/libs/vis_utils.py,sha256=hTOTEakKV0pGMbk0tbRkVI_tku7A3dGc96ynObZ4kwI,10207
 magic_pdf/model/360_layout_analysis.py,sha256=GbchKPJRVcrxvwNXMnR4vt8lOLPauTWMl-43ayyhX7U,221
 magic_pdf/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -117,8 +117,8 @@ magic_pdf/train_utils/convert_to_train_format.py,sha256=ifo2FAoBMa_etCvz0O4v03xO
 magic_pdf/train_utils/extract_caption.py,sha256=gommEqIEWLplSDEJWD7_66daqlOBsWhpRBW1DHpkny4,1825
 magic_pdf/train_utils/remove_footer_header.py,sha256=pyeNNdJ-th3wl5Xwb10ZLYNaFN4-6BmahoMFE8VTNNs,5978
 magic_pdf/train_utils/vis_utils.py,sha256=MV9N9cT3ifJ35u7LFKGF9I_bOIQrtU1zcsxu2hj3aqM,10111
-magic_pdf-0.5.7.dist-info/LICENSE.md,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-magic_pdf-0.5.7.dist-info/METADATA,sha256=6tyRzBGDgaq7hCfgzI_KjOhnW_nStuIse-6bmB8WxN8,814
-magic_pdf-0.5.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-magic_pdf-0.5.7.dist-info/top_level.txt,sha256=J9I0AzmHWGkp9c6DL8Oe4mEx3yYphLzkRn4H25Lg1rE,10
-magic_pdf-0.5.7.dist-info/RECORD,,
+magic_pdf-0.5.8.dist-info/LICENSE.md,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+magic_pdf-0.5.8.dist-info/METADATA,sha256=Z7HrhP7T0_dQOjCX-CztMe77Mbt90IoY8JH0IhmRHH0,814
+magic_pdf-0.5.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+magic_pdf-0.5.8.dist-info/top_level.txt,sha256=J9I0AzmHWGkp9c6DL8Oe4mEx3yYphLzkRn4H25Lg1rE,10
+magic_pdf-0.5.8.dist-info/RECORD,,