api.rst 3.69 KB
Newer Older
xu rui's avatar
xu rui committed
1

xu rui's avatar
xu rui committed
2
3
Api Usage 
===========
icecraft's avatar
icecraft committed
4
5
6
7

Local File Example
^^^^^^^^^^^^^^^^^^

xu rui's avatar
xu rui committed
8
9
10
11
12
.. code:: python

    import os

    from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
xu rui's avatar
xu rui committed
13
14
    from magic_pdf.data.dataset import PymuDocDataset
    from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
xu rui's avatar
xu rui committed
15

xu rui's avatar
xu rui committed
16
    # args
xu rui's avatar
xu rui committed
17
    pdf_file_name = "abc.pdf"  # replace with the real pdf path
xu rui's avatar
xu rui committed
18
    name_without_suff = pdf_file_name.split(".")[0]
xu rui's avatar
xu rui committed
19

xu rui's avatar
xu rui committed
20
    # prepare env
xu rui's avatar
xu rui committed
21
    local_image_dir, local_md_dir = "output/images", "output"
xu rui's avatar
xu rui committed
22
23
    image_dir = str(os.path.basename(local_image_dir))

xu rui's avatar
xu rui committed
24
25
26
27
    os.makedirs(local_image_dir, exist_ok=True)

    image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(
        local_md_dir
icecraft's avatar
icecraft committed
28
    )
xu rui's avatar
xu rui committed
29
30
    image_dir = str(os.path.basename(local_image_dir))

xu rui's avatar
xu rui committed
31
    # read bytes
xu rui's avatar
xu rui committed
32
    reader1 = FileBasedDataReader("")
xu rui's avatar
xu rui committed
33
    pdf_bytes = reader1.read(pdf_file_name)  # read the pdf content
xu rui's avatar
xu rui committed
34

xu rui's avatar
xu rui committed
35
36
37
    # proc
    ## Create Dataset Instance
    ds = PymuDocDataset(pdf_bytes)
xu rui's avatar
xu rui committed
38

xu rui's avatar
xu rui committed
39
40
    ## inference 
    infer_result = ds.apply(doc_analyze, ocr=True)
xu rui's avatar
xu rui committed
41

xu rui's avatar
xu rui committed
42
43
    ### draw model result on each page
    infer_result.draw_model(os.path.join(local_md_dir, f"{name_without_suff}_model.pdf"))
xu rui's avatar
xu rui committed
44

xu rui's avatar
xu rui committed
45
46
    ## pipeline
    pipe_result = infer_result.pipe_ocr_mode(image_writer)
xu rui's avatar
xu rui committed
47

xu rui's avatar
xu rui committed
48
49
    ### draw layout result on each page
    pipe_result.draw_layout(os.path.join(local_md_dir, f"{name_without_suff}_layout.pdf"))
xu rui's avatar
xu rui committed
50

xu rui's avatar
xu rui committed
51
52
    ### draw spans result on each page
    pipe_result.draw_span(os.path.join(local_md_dir, f"{name_without_suff}_spans.pdf"))
xu rui's avatar
xu rui committed
53

xu rui's avatar
xu rui committed
54
55
    ### dump markdown
    pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir)
xu rui's avatar
xu rui committed
56
57


icecraft's avatar
icecraft committed
58
59
60
61
62
63
64
65
S3 File Example
^^^^^^^^^^^^^^^^

.. code:: python

    import os

    from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter
xu rui's avatar
xu rui committed
66
67
    from magic_pdf.data.dataset import PymuDocDataset
    from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
icecraft's avatar
icecraft committed
68
69
70
71
72
73
74
75
76
77
78

    bucket_name = "{Your S3 Bucket Name}"  # replace with real bucket name
    ak = "{Your S3 access key}"  # replace with real s3 access key
    sk = "{Your S3 secret key}"  # replace with real s3 secret key
    endpoint_url = "{Your S3 endpoint_url}"  # replace with real s3 endpoint_url


    reader = S3DataReader('unittest/tmp/', bucket_name, ak, sk, endpoint_url)  # replace `unittest/tmp` with the real s3 prefix
    writer = S3DataWriter('unittest/tmp', bucket_name, ak, sk, endpoint_url)
    image_writer = S3DataWriter('unittest/tmp/images', bucket_name, ak, sk, endpoint_url)

xu rui's avatar
xu rui committed
79
80
81
82
83
84
85
86
    # args
    pdf_file_name = (
        "s3://llm-pdf-text-1/unittest/tmp/bug5-11.pdf"  # replace with the real s3 path
    )

    # prepare env
    local_dir = "output"
    name_without_suff = os.path.basename(pdf_file_name).split(".")[0]
icecraft's avatar
icecraft committed
87

xu rui's avatar
xu rui committed
88
    # read bytes
icecraft's avatar
icecraft committed
89
90
    pdf_bytes = reader.read(pdf_file_name)  # read the pdf content

xu rui's avatar
xu rui committed
91
92
93
    # proc
    ## Create Dataset Instance
    ds = PymuDocDataset(pdf_bytes)
icecraft's avatar
icecraft committed
94

xu rui's avatar
xu rui committed
95
96
    ## inference 
    infer_result = ds.apply(doc_analyze, ocr=True)
icecraft's avatar
icecraft committed
97

xu rui's avatar
xu rui committed
98
99
    ### draw model result on each page
    infer_result.draw_model(os.path.join(local_dir, f'{name_without_suff}_model.pdf'))  # dump to local
icecraft's avatar
icecraft committed
100

xu rui's avatar
xu rui committed
101
102
    ## pipeline
    pipe_result = infer_result.pipe_ocr_mode(image_writer)
icecraft's avatar
icecraft committed
103

xu rui's avatar
xu rui committed
104
105
106
107
108
    ### draw layout result on each page
    pipe_result.draw_layout(os.path.join(local_dir, f'{name_without_suff}_layout.pdf'))  # dump to local

    ### draw spans result on each page
    pipe_result.draw_span(os.path.join(local_dir, f'{name_without_suff}_spans.pdf'))   # dump to local 
icecraft's avatar
icecraft committed
109

xu rui's avatar
xu rui committed
110
111
    ### dump markdown
    pipe_result.dump_md(writer, f'{name_without_suff}.md', "unittest/tmp/images")    # dump to remote s3
icecraft's avatar
icecraft committed
112
113


xu rui's avatar
xu rui committed
114
Check :doc:`../data/data_reader_writer` for more [reader | writer] examples and check :doc:`../../api/pipe_operators` or :doc:`../../api/model_operators` for api details