magic-pdf 0.9.2__py3-none-any.whl → 0.10.0__py3-none-any.whl

magic_pdf/pipe/UNIPipe.py

@@ -2,22 +2,21 @@ import json
 
 from loguru import logger
 
-from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
-from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
-from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
-from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
+from magic_pdf.config.make_content_config import DropMode, MakeMode
+from magic_pdf.data.data_reader_writer import DataWriter
 from magic_pdf.libs.commons import join_path
+from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
 from magic_pdf.pipe.AbsPipe import AbsPipe
-from magic_pdf.user_api import parse_union_pdf, parse_ocr_pdf
+from magic_pdf.user_api import parse_ocr_pdf, parse_union_pdf
 
 
 class UNIPipe(AbsPipe):
 
-    def __init__(self, pdf_bytes: bytes, jso_useful_key: dict, image_writer: AbsReaderWriter, is_debug: bool = False,
+    def __init__(self, pdf_bytes: bytes, jso_useful_key: dict, image_writer: DataWriter, is_debug: bool = False,
                  start_page_id=0, end_page_id=None, lang=None,
                  layout_model=None, formula_enable=None, table_enable=None):
-        self.pdf_type = jso_useful_key["_pdf_type"]
-        super().__init__(pdf_bytes, jso_useful_key["model_list"], image_writer, is_debug, start_page_id, end_page_id,
+        self.pdf_type = jso_useful_key['_pdf_type']
+        super().__init__(pdf_bytes, jso_useful_key['model_list'], image_writer, is_debug, start_page_id, end_page_id,
                          lang, layout_model, formula_enable, table_enable)
         if len(self.model_list) == 0:
             self.input_model_is_empty = True
@@ -54,27 +53,28 @@ class UNIPipe(AbsPipe):
 
     def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.NONE_WITH_REASON):
         result = super().pipe_mk_uni_format(img_parent_path, drop_mode)
-        logger.info("uni_pipe mk content list finished")
+        logger.info('uni_pipe mk content list finished')
         return result
 
     def pipe_mk_markdown(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD):
         result = super().pipe_mk_markdown(img_parent_path, drop_mode, md_make_mode)
-        logger.info(f"uni_pipe mk {md_make_mode} finished")
+        logger.info(f'uni_pipe mk {md_make_mode} finished')
         return result
 
 
 if __name__ == '__main__':
     # 测试
-    drw = DiskReaderWriter(r"D:/project/20231108code-clean")
+    from magic_pdf.data.data_reader_writer import DataReader
+    drw = DataReader(r'D:/project/20231108code-clean')
 
-    pdf_file_path = r"linshixuqiu\19983-00.pdf"
-    model_file_path = r"linshixuqiu\19983-00.json"
-    pdf_bytes = drw.read(pdf_file_path, AbsReaderWriter.MODE_BIN)
-    model_json_txt = drw.read(model_file_path, AbsReaderWriter.MODE_TXT)
+    pdf_file_path = r'linshixuqiu\19983-00.pdf'
+    model_file_path = r'linshixuqiu\19983-00.json'
+    pdf_bytes = drw.read(pdf_file_path)
+    model_json_txt = drw.read(model_file_path).decode()
     model_list = json.loads(model_json_txt)
-    write_path = r"D:\project\20231108code-clean\linshixuqiu\19983-00"
-    img_bucket_path = "imgs"
-    img_writer = DiskReaderWriter(join_path(write_path, img_bucket_path))
+    write_path = r'D:\project\20231108code-clean\linshixuqiu\19983-00'
+    img_bucket_path = 'imgs'
+    img_writer = DataWriter(join_path(write_path, img_bucket_path))
 
     # pdf_type = UNIPipe.classify(pdf_bytes)
     # jso_useful_key = {
@@ -83,8 +83,8 @@ if __name__ == '__main__':
     # }
 
     jso_useful_key = {
-        "_pdf_type": "",
-        "model_list": model_list
+        '_pdf_type': '',
+        'model_list': model_list
     }
     pipe = UNIPipe(pdf_bytes, jso_useful_key, img_writer)
     pipe.pipe_classify()
@@ -92,8 +92,7 @@ if __name__ == '__main__':
     md_content = pipe.pipe_mk_markdown(img_bucket_path)
     content_list = pipe.pipe_mk_uni_format(img_bucket_path)
 
-    md_writer = DiskReaderWriter(write_path)
-    md_writer.write(md_content, "19983-00.md", AbsReaderWriter.MODE_TXT)
-    md_writer.write(json.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4), "19983-00.json",
-                    AbsReaderWriter.MODE_TXT)
-    md_writer.write(str(content_list), "19983-00.txt", AbsReaderWriter.MODE_TXT)
+    md_writer = DataWriter(write_path)
+    md_writer.write_string('19983-00.md', md_content)
+    md_writer.write_string('19983-00.json', json.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4))
+    md_writer.write_string('19983-00.txt', str(content_list))

magic_pdf/post_proc/pdf_post_filter.py

@@ -1,19 +1,17 @@
 from loguru import logger
 
+from magic_pdf.config.drop_reason import DropReason
 from magic_pdf.layout.layout_sort import get_columns_cnt_of_layout
-from magic_pdf.libs.drop_reason import DropReason
 
 
 def __is_pseudo_single_column(page_info) -> bool:
-    """
-    判断一个页面是否伪单列。
+    """判断一个页面是否伪单列。
 
     Args:
         page_info (dict): 页面信息字典，包括'_layout_tree'和'preproc_blocks'。
 
     Returns:
         Tuple[bool, Optional[str]]: 如果页面伪单列返回(True, extra_info)，否则返回(False, None)。
-
     """
     layout_tree = page_info['_layout_tree']
     layout_column_width = get_columns_cnt_of_layout(layout_tree)
@@ -41,27 +39,22 @@ def __is_pseudo_single_column(page_info) -> bool:
             if num_lines > 20:
                 radio = num_satisfying_lines / num_lines
                 if radio >= 0.5:
-                    extra_info = f"{{num_lines: {num_lines}, num_satisfying_lines: {num_satisfying_lines}}}"
+                    extra_info = f'{{num_lines: {num_lines}, num_satisfying_lines: {num_satisfying_lines}}}'
                     block_text = []
                     for line in lines:
                         if line['spans']:
                             for span in line['spans']:
                                 block_text.append(span['text'])
-                    logger.warning(f"pseudo_single_column block_text: {block_text}")
+                    logger.warning(f'pseudo_single_column block_text: {block_text}')
                     return True, extra_info
 
     return False, None
 
 
 def pdf_post_filter(page_info) -> tuple:
-    """
-    return:(True|False, err_msg)
-        True, 如果pdf符合要求
-        False, 如果pdf不符合要求
-
-    """
+    """return:(True|False, err_msg) True, 如果pdf符合要求 False, 如果pdf不符合要求."""
     bool_is_pseudo_single_column, extra_info = __is_pseudo_single_column(page_info)
     if bool_is_pseudo_single_column:
-        return False, {"_need_drop": True, "_drop_reason": DropReason.PSEUDO_SINGLE_COLUMN, "extra_info": extra_info}
+        return False, {'_need_drop': True, '_drop_reason': DropReason.PSEUDO_SINGLE_COLUMN, 'extra_info': extra_info}
 
-    return True, None
+    return True, None

magic_pdf/pre_proc/cut_image.py

@@ -1,7 +1,7 @@
 from loguru import logger
 
+from magic_pdf.config.ocr_content_type import ContentType
 from magic_pdf.libs.commons import join_path
-from magic_pdf.libs.ocr_content_type import ContentType
 from magic_pdf.libs.pdf_image_tools import cut_image
 
 
@@ -29,9 +29,7 @@ def txt_save_images_by_bboxes(page_num: int, page, pdf_bytes_md5: str,
                               image_bboxes: list, images_overlap_backup: list, table_bboxes: list,
                               equation_inline_bboxes: list,
                               equation_interline_bboxes: list, imageWriter) -> dict:
-    """
-    返回一个dict, key为bbox, 值是图片地址
-    """
+    """返回一个dict, key为bbox, 值是图片地址."""
     image_info = []
     image_backup_info = []
     table_info = []
@@ -46,26 +44,26 @@ def txt_save_images_by_bboxes(page_num: int, page, pdf_bytes_md5: str,
     for bbox in image_bboxes:
         if not check_img_bbox(bbox):
             continue
-        image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
-        image_info.append({"bbox": bbox, "image_path": image_path})
+        image_path = cut_image(bbox, page_num, page, return_path('images'), imageWriter)
+        image_info.append({'bbox': bbox, 'image_path': image_path})
 
     for bbox in images_overlap_backup:
         if not check_img_bbox(bbox):
             continue
-        image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
-        image_backup_info.append({"bbox": bbox, "image_path": image_path})
+        image_path = cut_image(bbox, page_num, page, return_path('images'), imageWriter)
+        image_backup_info.append({'bbox': bbox, 'image_path': image_path})
 
     for bbox in table_bboxes:
         if not check_img_bbox(bbox):
             continue
-        image_path = cut_image(bbox, page_num, page, return_path("tables"), imageWriter)
-        table_info.append({"bbox": bbox, "image_path": image_path})
+        image_path = cut_image(bbox, page_num, page, return_path('tables'), imageWriter)
+        table_info.append({'bbox': bbox, 'image_path': image_path})
 
     return image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info
 
 
 def check_img_bbox(bbox) -> bool:
     if any([bbox[0] >= bbox[2], bbox[1] >= bbox[3]]):
-        logger.warning(f"image_bboxes: 错误的box, {bbox}")
+        logger.warning(f'image_bboxes: 错误的box, {bbox}')
         return False
     return True