magic-pdf 0.6.1__py3-none-any.whl → 0.7.0a1__py3-none-any.whl

magic_pdf/resources/model_config/layoutlmv3/layoutlmv3_base_inference.yaml

@@ -1,6 +1,6 @@
 AUG:
   DETR: true
-CACHE_DIR: /mnt/localdata/users/yupanhuang/cache/huggingface
+CACHE_DIR: ~/cache/huggingface
 CUDNN_BENCHMARK: false
 DATALOADER:
   ASPECT_RATIO_GROUPING: true
@@ -294,7 +294,7 @@ MODEL:
     POS_TYPE: abs
   WEIGHTS: 
 OUTPUT_DIR: 
-SCIHUB_DATA_DIR_TRAIN: /mnt/petrelfs/share_data/zhaozhiyuan/publaynet/layout_scihub/train
+SCIHUB_DATA_DIR_TRAIN: ~/publaynet/layout_scihub/train
 SEED: 42
 SOLVER:
   AMP:

magic_pdf/resources/model_config/model_configs.yaml

@@ -2,8 +2,12 @@ config:
   device: cpu
   layout: True
   formula: True
+  table_config:
+    is_table_recog_enable: False
+    max_time: 400
 
 weights:
   layout: Layout/model_final.pth
   mfd: MFD/weights.pt
   mfr: MFR/UniMERNet
+  table: TabRec/StructEqTable

magic_pdf/rw/AbsReaderWriter.py

@@ -2,33 +2,16 @@ from abc import ABC, abstractmethod
 
 
 class AbsReaderWriter(ABC):
-    """
-    同时支持二进制和文本读写的抽象类
-    """
     MODE_TXT = "text"
     MODE_BIN = "binary"
-
-    def __init__(self, parent_path):
-        # 初始化代码可以在这里添加，如果需要的话
-        self.parent_path = parent_path # 对于本地目录是父目录，对于s3是会写到这个path下。
-
     @abstractmethod
     def read(self, path: str, mode=MODE_TXT):
-        """
-        无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path
-        """
         raise NotImplementedError
 
     @abstractmethod
     def write(self, content: str, path: str, mode=MODE_TXT):
-        """
-        无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path
-        """
         raise NotImplementedError
 
     @abstractmethod
-    def read_jsonl(self, path: str, byte_start=0, byte_end=None, encoding='utf-8'):
-        """
-        无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path
-        """
+    def read_offset(self, path: str, offset=0, limit=None) -> bytes:
         raise NotImplementedError

magic_pdf/rw/DiskReaderWriter.py

@@ -3,34 +3,29 @@ from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
 from loguru import logger
 
 
-MODE_TXT = "text"
-MODE_BIN = "binary"
-
-
 class DiskReaderWriter(AbsReaderWriter):
-
     def __init__(self, parent_path, encoding="utf-8"):
         self.path = parent_path
         self.encoding = encoding
 
-    def read(self, path, mode=MODE_TXT):
+    def read(self, path, mode=AbsReaderWriter.MODE_TXT):
         if os.path.isabs(path):
             abspath = path
         else:
             abspath = os.path.join(self.path, path)
         if not os.path.exists(abspath):
-            logger.error(f"文件 {abspath} 不存在")
-            raise Exception(f"文件 {abspath} 不存在")
-        if mode == MODE_TXT:
+            logger.error(f"file {abspath} not exists")
+            raise Exception(f"file {abspath} no exists")
+        if mode == AbsReaderWriter.MODE_TXT:
             with open(abspath, "r", encoding=self.encoding) as f:
                 return f.read()
-        elif mode == MODE_BIN:
+        elif mode == AbsReaderWriter.MODE_BIN:
             with open(abspath, "rb") as f:
                 return f.read()
         else:
             raise ValueError("Invalid mode. Use 'text' or 'binary'.")
 
-    def write(self, content, path, mode=MODE_TXT):
+    def write(self, content, path, mode=AbsReaderWriter.MODE_TXT):
         if os.path.isabs(path):
             abspath = path
         else:
@@ -38,29 +33,42 @@ class DiskReaderWriter(AbsReaderWriter):
         directory_path = os.path.dirname(abspath)
         if not os.path.exists(directory_path):
             os.makedirs(directory_path)
-        if mode == MODE_TXT:
+        if mode == AbsReaderWriter.MODE_TXT:
             with open(abspath, "w", encoding=self.encoding, errors="replace") as f:
                 f.write(content)
 
-        elif mode == MODE_BIN:
+        elif mode == AbsReaderWriter.MODE_BIN:
             with open(abspath, "wb") as f:
                 f.write(content)
         else:
             raise ValueError("Invalid mode. Use 'text' or 'binary'.")
 
-    def read_jsonl(self, path: str, byte_start=0, byte_end=None, encoding="utf-8"):
-        return self.read(path)
+    def read_offset(self, path: str, offset=0, limit=None):
+        abspath = path
+        if not os.path.isabs(path):
+            abspath = os.path.join(self.path, path)
+        with open(abspath, "rb") as f:
+            f.seek(offset)
+            return f.read(limit)
 
 
-# 使用示例
 if __name__ == "__main__":
-    file_path = "io/test/example.txt"
-    drw = DiskReaderWriter("D:\projects\papayfork\Magic-PDF\magic_pdf")
+    if 0:
+        file_path = "io/test/example.txt"
+        drw = DiskReaderWriter("D:\projects\papayfork\Magic-PDF\magic_pdf")
+
+        # 写入内容到文件
+        drw.write(b"Hello, World!", path="io/test/example.txt", mode="binary")
+
+        # 从文件读取内容
+        content = drw.read(path=file_path)
+        if content:
+            logger.info(f"从 {file_path} 读取的内容: {content}")
+    if 1:
+        drw = DiskReaderWriter("/opt/data/pdf/resources/test/io/")
+        content_bin = drw.read_offset("1.txt")
+        assert content_bin == b"ABCD!"
 
-    # 写入内容到文件
-    drw.write(b"Hello, World!", path="io/test/example.txt", mode="binary")
+        content_bin = drw.read_offset("1.txt", offset=1, limit=2)
+        assert content_bin == b"BC"
 
-    # 从文件读取内容
-    content = drw.read(path=file_path)
-    if content:
-        logger.info(f"从 {file_path} 读取的内容: {content}")

magic_pdf/rw/S3ReaderWriter.py

@@ -2,16 +2,18 @@ from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
 from magic_pdf.libs.commons import parse_aws_param, parse_bucket_key, join_path
 import boto3
 from loguru import logger
-from boto3.s3.transfer import TransferConfig
 from botocore.config import Config
-import os
-
-MODE_TXT = "text"
-MODE_BIN = "binary"
 
 
 class S3ReaderWriter(AbsReaderWriter):
-    def __init__(self, ak: str, sk: str, endpoint_url: str, addressing_style: str = 'auto', parent_path: str = ''):
+    def __init__(
+        self,
+        ak: str,
+        sk: str,
+        endpoint_url: str,
+        addressing_style: str = "auto",
+        parent_path: str = "",
+    ):
         self.client = self._get_client(ak, sk, endpoint_url, addressing_style)
         self.path = parent_path
 
@@ -21,12 +23,14 @@ class S3ReaderWriter(AbsReaderWriter):
             aws_access_key_id=ak,
             aws_secret_access_key=sk,
             endpoint_url=endpoint_url,
-            config=Config(s3={"addressing_style": addressing_style},
-                          retries={'max_attempts': 5, 'mode': 'standard'}),
+            config=Config(
+                s3={"addressing_style": addressing_style},
+                retries={"max_attempts": 5, "mode": "standard"},
+            ),
         )
         return s3_client
 
-    def read(self, s3_relative_path, mode=MODE_TXT, encoding="utf-8"):
+    def read(self, s3_relative_path, mode=AbsReaderWriter.MODE_TXT, encoding="utf-8"):
         if s3_relative_path.startswith("s3://"):
             s3_path = s3_relative_path
         else:
@@ -34,22 +38,22 @@ class S3ReaderWriter(AbsReaderWriter):
         bucket_name, key = parse_bucket_key(s3_path)
         res = self.client.get_object(Bucket=bucket_name, Key=key)
         body = res["Body"].read()
-        if mode == MODE_TXT:
+        if mode == AbsReaderWriter.MODE_TXT:
             data = body.decode(encoding)  # Decode bytes to text
-        elif mode == MODE_BIN:
+        elif mode == AbsReaderWriter.MODE_BIN:
             data = body
         else:
             raise ValueError("Invalid mode. Use 'text' or 'binary'.")
         return data
 
-    def write(self, content, s3_relative_path, mode=MODE_TXT, encoding="utf-8"):
+    def write(self, content, s3_relative_path, mode=AbsReaderWriter.MODE_TXT, encoding="utf-8"):
         if s3_relative_path.startswith("s3://"):
             s3_path = s3_relative_path
         else:
             s3_path = join_path(self.path, s3_relative_path)
-        if mode == MODE_TXT:
+        if mode == AbsReaderWriter.MODE_TXT:
             body = content.encode(encoding)  # Encode text data as bytes
-        elif mode == MODE_BIN:
+        elif mode == AbsReaderWriter.MODE_BIN:
             body = content
         else:
             raise ValueError("Invalid mode. Use 'text' or 'binary'.")
@@ -57,51 +61,82 @@ class S3ReaderWriter(AbsReaderWriter):
         self.client.put_object(Body=body, Bucket=bucket_name, Key=key)
         logger.info(f"内容已写入 {s3_path} ")
 
-    def read_jsonl(self, path: str, byte_start=0, byte_end=None, mode=MODE_TXT, encoding='utf-8'):
+    def read_offset(self, path: str, offset=0, limit=None) -> bytes:
         if path.startswith("s3://"):
             s3_path = path
         else:
             s3_path = join_path(self.path, path)
         bucket_name, key = parse_bucket_key(s3_path)
 
-        range_header = f'bytes={byte_start}-{byte_end}' if byte_end else f'bytes={byte_start}-'
+        range_header = (
+            f"bytes={offset}-{offset+limit-1}" if limit else f"bytes={offset}-"
+        )
         res = self.client.get_object(Bucket=bucket_name, Key=key, Range=range_header)
-        body = res["Body"].read()
-        if mode == MODE_TXT:
-            data = body.decode(encoding)  # Decode bytes to text
-        elif mode == MODE_BIN:
-            data = body
-        else:
-            raise ValueError("Invalid mode. Use 'text' or 'binary'.")
-        return data
+        return res["Body"].read()
 
 
 if __name__ == "__main__":
-    # Config the connection info
-    ak = ""
-    sk = ""
-    endpoint_url = ""
-    addressing_style = "auto"
-    bucket_name = ""
-    # Create an S3ReaderWriter object
-    s3_reader_writer = S3ReaderWriter(ak, sk, endpoint_url, addressing_style, "s3://bucket_name/")
+    if 0:
+        # Config the connection info
+        ak = ""
+        sk = ""
+        endpoint_url = ""
+        addressing_style = "auto"
+        bucket_name = ""
+        # Create an S3ReaderWriter object
+        s3_reader_writer = S3ReaderWriter(
+            ak, sk, endpoint_url, addressing_style, "s3://bucket_name/"
+        )
 
-    # Write text data to S3
-    text_data = "This is some text data"
-    s3_reader_writer.write(data=text_data, s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_TXT)
+        # Write text data to S3
+        text_data = "This is some text data"
+        s3_reader_writer.write(
+            text_data,
+            s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json",
+            mode=AbsReaderWriter.MODE_TXT,
+        )
+
+        # Read text data from S3
+        text_data_read = s3_reader_writer.read(
+            s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=AbsReaderWriter.MODE_TXT
+        )
+        logger.info(f"Read text data from S3: {text_data_read}")
+        # Write binary data to S3
+        binary_data = b"This is some binary data"
+        s3_reader_writer.write(
+            text_data,
+            s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json",
+            mode=AbsReaderWriter.MODE_BIN,
+        )
 
-    # Read text data from S3
-    text_data_read = s3_reader_writer.read(s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_TXT)
-    logger.info(f"Read text data from S3: {text_data_read}")
-    # Write binary data to S3
-    binary_data = b"This is some binary data"
-    s3_reader_writer.write(data=text_data, s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_BIN)
+        # Read binary data from S3
+        binary_data_read = s3_reader_writer.read(
+            s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=AbsReaderWriter.MODE_BIN
+        )
+        logger.info(f"Read binary data from S3: {binary_data_read}")
+
+        # Range Read text data from S3
+        binary_data_read = s3_reader_writer.read_offset(
+            path=f"s3://{bucket_name}/ebook/test/test.json", offset=0, limit=10
+        )
+        logger.info(f"Read binary data from S3: {binary_data_read}")
+    if 1:
+        import os
+        import json
 
-    # Read binary data from S3
-    binary_data_read = s3_reader_writer.read(s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_BIN)
-    logger.info(f"Read binary data from S3: {binary_data_read}")
+        ak = os.getenv("AK", "")
+        sk = os.getenv("SK", "")
+        endpoint_url = os.getenv("ENDPOINT", "")
+        bucket = os.getenv("S3_BUCKET", "")
+        prefix = os.getenv("S3_PREFIX", "")
+        key_basename = os.getenv("S3_KEY_BASENAME", "")
+        s3_reader_writer = S3ReaderWriter(
+            ak, sk, endpoint_url, "auto", f"s3://{bucket}/{prefix}"
+        )
+        content_bin = s3_reader_writer.read_offset(key_basename)
+        assert content_bin[:10] == b'{"track_id'
+        assert content_bin[-10:] == b'r":null}}\n'
 
-    # Range Read text data from S3
-    binary_data_read = s3_reader_writer.read_jsonl(path=f"s3://{bucket_name}/ebook/test/test.json",
-                                                   byte_start=0, byte_end=10, mode=MODE_BIN)
-    logger.info(f"Read binary data from S3: {binary_data_read}")
+        content_bin = s3_reader_writer.read_offset(key_basename, offset=424, limit=426)
+        jso = json.dumps(content_bin.decode("utf-8"))
+        print(jso)

magic_pdf/tools/cli.py

@@ -0,0 +1,79 @@
+import os
+import click
+from loguru import logger
+from pathlib import Path
+
+from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
+from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
+import magic_pdf.model as model_config
+from magic_pdf.tools.common import parse_pdf_methods, do_parse
+from magic_pdf.libs.version import __version__
+
+
+@click.command()
+@click.version_option(__version__, "--version", "-v", help="display the version and exit")
+@click.option(
+    "-p",
+    "--path",
+    "path",
+    type=click.Path(exists=True),
+    required=True,
+    help="local pdf filepath or directory",
+)
+@click.option(
+    "-o",
+    "--output-dir",
+    "output_dir",
+    type=str,
+    help="output local directory",
+    default="",
+)
+@click.option(
+    "-m",
+    "--method",
+    "method",
+    type=parse_pdf_methods,
+    help="""the method for parsing pdf. 
+ocr: using ocr technique to extract information from pdf.
+txt: suitable for the text-based pdf only and outperform ocr.
+auto: automatically choose the best method for parsing pdf from ocr and txt.
+without method specified, auto will be used by default.""",
+    default="auto",
+)
+def cli(path, output_dir, method):
+    model_config.__use_inside_model__ = True
+    model_config.__model_mode__ = "full"
+    if output_dir == "":
+        if os.path.isdir(path):
+            output_dir = os.path.join(path, "output")
+        else:
+            output_dir = os.path.join(os.path.dirname(path), "output")
+
+    def read_fn(path):
+        disk_rw = DiskReaderWriter(os.path.dirname(path))
+        return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
+
+    def parse_doc(doc_path: str):
+        try:
+            file_name = str(Path(doc_path).stem)
+            pdf_data = read_fn(doc_path)
+            do_parse(
+                output_dir,
+                file_name,
+                pdf_data,
+                [],
+                method,
+            )
+
+        except Exception as e:
+            logger.exception(e)
+
+    if os.path.isdir(path):
+        for doc_path in Path(path).glob("*.pdf"):
+            parse_doc(doc_path)
+    else:
+        parse_doc(path)
+
+
+if __name__ == "__main__":
+    cli()

magic_pdf/tools/cli_dev.py

@@ -0,0 +1,156 @@
+import os
+import json as json_parse
+import click
+from pathlib import Path
+from magic_pdf.libs.path_utils import (
+    parse_s3path,
+    parse_s3_range_params,
+    remove_non_official_s3_args,
+)
+from magic_pdf.libs.config_reader import (
+    get_s3_config,
+)
+from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter
+from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
+from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
+import magic_pdf.model as model_config
+from magic_pdf.tools.common import parse_pdf_methods, do_parse
+from magic_pdf.libs.version import __version__
+
+
+def read_s3_path(s3path):
+    bucket, key = parse_s3path(s3path)
+
+    s3_ak, s3_sk, s3_endpoint = get_s3_config(bucket)
+    s3_rw = S3ReaderWriter(
+        s3_ak, s3_sk, s3_endpoint, "auto", remove_non_official_s3_args(s3path)
+    )
+    may_range_params = parse_s3_range_params(s3path)
+    if may_range_params is None or 2 != len(may_range_params):
+        byte_start, byte_end = 0, None
+    else:
+        byte_start, byte_end = int(may_range_params[0]), int(may_range_params[1])
+        byte_end += byte_start - 1
+    return s3_rw.read_jsonl(
+        remove_non_official_s3_args(s3path),
+        byte_start,
+        byte_end,
+        AbsReaderWriter.MODE_BIN,
+    )
+
+
+@click.group()
+@click.version_option(__version__, "--version", "-v", help="显示版本信息")
+def cli():
+    pass
+
+
+@cli.command()
+@click.option(
+    "-j",
+    "--jsonl",
+    "jsonl",
+    type=str,
+    help="输入 jsonl 路径，本地或者 s3 上的文件",
+    required=True,
+)
+@click.option(
+    "-m",
+    "--method",
+    "method",
+    type=parse_pdf_methods,
+    help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
+    default="auto",
+)
+@click.option(
+    "-o",
+    "--output-dir",
+    "output_dir",
+    type=str,
+    help="输出到本地目录",
+    default="",
+)
+def jsonl(jsonl, method, output_dir):
+    print("haha")
+    model_config.__use_inside_model__ = False
+    full_jsonl_path = os.path.realpath(jsonl)
+    if output_dir == "":
+        output_dir = os.path.join(os.path.dirname(full_jsonl_path), "output")
+
+    if jsonl.startswith("s3://"):
+        jso = json_parse.loads(read_s3_path(jsonl).decode("utf-8"))
+    else:
+        with open(jsonl) as f:
+            jso = json_parse.loads(f.readline())
+    s3_file_path = jso.get("file_location")
+    if s3_file_path is None:
+        s3_file_path = jso.get("path")
+    pdf_file_name = Path(s3_file_path).stem
+    pdf_data = read_s3_path(s3_file_path)
+
+
+    print(pdf_file_name, jso, method)
+    do_parse(
+        output_dir,
+        pdf_file_name,
+        pdf_data,
+        jso["doc_layout_result"],
+        method,
+        f_dump_content_list=True,
+    )
+
+
+@cli.command()
+@click.option(
+    "-p",
+    "--pdf",
+    "pdf",
+    type=click.Path(exists=True),
+    required=True,
+    help="本地 PDF 文件",
+)
+@click.option(
+    "-j",
+    "--json",
+    "json_data",
+    type=click.Path(exists=True),
+    required=True,
+    help="本地模型推理出的 json 数据",
+)
+@click.option(
+    "-o", "--output-dir", "output_dir", type=str, help="本地输出目录", default=""
+)
+@click.option(
+    "-m",
+    "--method",
+    "method",
+    type=parse_pdf_methods,
+    help="指定解析方法。txt: 文本型 pdf 解析方法， ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
+    default="auto",
+)
+def pdf(pdf, json_data, output_dir, method):
+    model_config.__use_inside_model__ = False
+    full_pdf_path = os.path.realpath(pdf)
+    if output_dir == "":
+        output_dir = os.path.join(os.path.dirname(full_pdf_path), "output")
+
+    def read_fn(path):
+        disk_rw = DiskReaderWriter(os.path.dirname(path))
+        return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
+
+    model_json_list = json_parse.loads(read_fn(json_data).decode("utf-8"))
+
+    file_name = str(Path(full_pdf_path).stem)
+    pdf_data = read_fn(full_pdf_path)
+    do_parse(
+        output_dir,
+        file_name,
+        pdf_data,
+        model_json_list,
+        method,
+        f_dump_content_list=True,
+    )
+
+
+if __name__ == "__main__":
+    cli()