test_cli_sdk.py 23.4 KB
Newer Older
yyy's avatar
yyy committed
1
2
3
4
5
6
"""test cli and sdk."""
import logging
import os
import pytest
from conf import conf
from lib import common
yyy's avatar
yyy committed
7
import time
yyy's avatar
yyy committed
8
import magic_pdf.model as model_config
dt-yy's avatar
dt-yy committed
9
10
11
from magic_pdf.data.read_api import read_local_images
from magic_pdf.pipe.UNIPipe import UNIPipe
from magic_pdf.data.read_api import read_local_office
dt-yy's avatar
dt-yy committed
12
13
from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter
from magic_pdf.config.make_content_config import DropMode, MakeMode
dt-yy's avatar
dt-yy committed
14
from magic_pdf.pipe.OCRPipe import OCRPipe
dt-yy's avatar
dt-yy committed
15
16
17
18
from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.config.enums import SupportedPdfParseMethod
yyy's avatar
yyy committed
19
20
21
22
model_config.__use_inside_model__ = True
pdf_res_path = conf.conf['pdf_res_path']
code_path = conf.conf['code_path']
pdf_dev_path = conf.conf['pdf_dev_path']
quyuan's avatar
quyuan committed
23
magic_pdf_config = "/home/quyuan/magic-pdf.json"
yyy's avatar
yyy committed
24
25
26

class TestCli:
    """test cli."""
quyuan's avatar
quyuan committed
27
28
29
30
31
32
33
34
35
    @pytest.fixture(autouse=True)
    def setup(self):
        """
        init
        """
        common.clear_gpu_memory()
        common.update_config_file(magic_pdf_config, "device-mode", "cuda")
        # 这里可以添加任何前置操作
        yield
yyy's avatar
yyy committed
36
37

    @pytest.mark.P0
dt-yy's avatar
dt-yy committed
38
    def test_pdf_local_sdk(self):
yyy's avatar
yyy committed
39
40
41
42
43
44
45
46
47
48
        """pdf sdk auto test."""
        demo_names = list()
        pdf_path = os.path.join(pdf_dev_path, 'pdf')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.pdf'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            pdf_path = os.path.join(pdf_dev_path, 'pdf', f'{demo_name}.pdf')
            local_image_dir = os.path.join(pdf_dev_path, 'pdf', 'images')
            image_dir = str(os.path.basename(local_image_dir))
dt-yy's avatar
dt-yy committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
            name_without_suff = os.path.basename(pdf_path).split(".pdf")[0]
            dir_path = os.path.join(pdf_dev_path, 'mineru')
            image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
            reader1 = FileBasedDataReader("")
            pdf_bytes = reader1.read(pdf_path)
            ds = PymuDocDataset(pdf_bytes)
            ## inference
            if ds.classify() == SupportedPdfParseMethod.OCR:
                infer_result = ds.apply(doc_analyze, ocr=True)
                ## pipeline
                pipe_result = infer_result.pipe_ocr_mode(image_writer)
            else:
                infer_result = ds.apply(doc_analyze, ocr=False)
                ## pipeline
                pipe_result = infer_result.pipe_txt_mode(image_writer)
            common.delete_file(dir_path)
            ### draw model result on each page
            infer_result.draw_model(os.path.join(dir_path, f"{name_without_suff}_model.pdf"))

            ### get model inference result
            model_inference_result = infer_result.get_infer_res()

            ### draw layout result on each page
            pipe_result.draw_layout(os.path.join(dir_path, f"{name_without_suff}_layout.pdf"))

            ### draw spans result on each page
            pipe_result.draw_span(os.path.join(dir_path, f"{name_without_suff}_spans.pdf"))

            ### dump markdown
            pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir)

            ### dump content list
            pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", image_dir)

            ### get markdown content
            md_content = pipe_result.get_markdown(image_dir, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD)

            ### get content list content
            content_list_content = pipe_result.get_content_list(image_dir, drop_mode=DropMode.NONE, md_make_mode=MakeMode.STANDARD_FORMAT)

            ### get middle json
            middle_json_content = pipe_result.get_middle_json()
            common.sdk_count_folders_and_check_contents(dir_path)

    @pytest.mark.P0
    def test_pdf_s3_sdk(self):
        """pdf s3 sdk test."""
        demo_names = list()
        pdf_path = os.path.join(pdf_dev_path, 'pdf')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.pdf'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            pdf_path = os.path.join(pdf_dev_path, 'pdf', f'{demo_name}.pdf')
            local_image_dir = os.path.join(pdf_dev_path, 'pdf', 'images')
            image_dir = str(os.path.basename(local_image_dir))
            name_without_suff = os.path.basename(pdf_path).split(".pdf")[0]
yyy's avatar
yyy committed
106
            dir_path = os.path.join(pdf_dev_path, 'mineru')
dt-yy's avatar
dt-yy committed
107
108
109
110
111
112
113
114
115
116
117
118
119
120
            image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
            reader1 = FileBasedDataReader("")
            pdf_bytes = reader1.read(pdf_path)
            ds = PymuDocDataset(pdf_bytes)
            ## inference
            if ds.classify() == SupportedPdfParseMethod.OCR:
                infer_result = ds.apply(doc_analyze, ocr=True)
                ## pipeline
                pipe_result = infer_result.pipe_ocr_mode(image_writer)
            else:
                infer_result = ds.apply(doc_analyze, ocr=False)
                ## pipeline
                pipe_result = infer_result.pipe_txt_mode(image_writer)
            common.delete_file(dir_path)
dt-yy's avatar
dt-yy committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
            ### draw model result on each page
            infer_result.draw_model(os.path.join(dir_path, f"{name_without_suff}_model.pdf"))

            ### get model inference result
            model_inference_result = infer_result.get_infer_res()

            ### draw layout result on each page
            pipe_result.draw_layout(os.path.join(dir_path, f"{name_without_suff}_layout.pdf"))

            ### draw spans result on each page
            pipe_result.draw_span(os.path.join(dir_path, f"{name_without_suff}_spans.pdf"))

            ### dump markdown
            pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir)

            ### dump content list
            pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", image_dir)

            ### get markdown content
            md_content = pipe_result.get_markdown(image_dir, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD)

            ### get content list content
            content_list_content = pipe_result.get_content_list(image_dir, drop_mode=DropMode.NONE, md_make_mode=MakeMode.STANDARD_FORMAT)

            ### get middle json
            middle_json_content = pipe_result.get_middle_json()
dt-yy's avatar
dt-yy committed
147
            common.sdk_count_folders_and_check_contents(dir_path)
yyy's avatar
yyy committed
148

dt-yy's avatar
dt-yy committed
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
    
    @pytest.mark.P0
    def test_pdf_local_ppt(self):
        """pdf sdk auto test."""
        demo_names = list()
        pdf_path = os.path.join(pdf_dev_path, 'ppt')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.pptx'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            pdf_path = os.path.join(pdf_dev_path, 'ppt', f'{demo_name}.pptx')
            local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
            image_dir = str(os.path.basename(local_image_dir))
            name_without_suff = os.path.basename(pdf_path).split(".pptx")[0]
            dir_path = os.path.join(pdf_dev_path, 'mineru')
            image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
            ds = read_local_office(pdf_path)[0]
            common.delete_file(dir_path)
            
            ds.apply(doc_analyze, ocr=True).pipe_txt_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir)          
            common.sdk_count_folders_and_check_contents(dir_path)



    @pytest.mark.P0
    def test_pdf_local_image(self):
        """pdf sdk auto test."""
        demo_names = list()
        pdf_path = os.path.join(pdf_dev_path, 'images')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.jpg'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            pdf_path = os.path.join(pdf_dev_path, 'images', f'{demo_name}.jpg')
            local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
            image_dir = str(os.path.basename(local_image_dir))
            name_without_suff = os.path.basename(pdf_path).split(".jpg")[0]
            dir_path = os.path.join(pdf_dev_path, 'mineru')
            common.delete_file(dir_path)
            image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
            ds = read_local_images(pdf_path)[0]
            ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(
            md_writer, f"{name_without_suff}.md", image_dir)
            common.sdk_count_folders_and_check_contents(dir_path)


    @pytest.mark.P0
    def test_local_image_dir(self):
        """local image dir."""
        demo_names = list()
        pdf_path = os.path.join(pdf_dev_path, 'images')
        dir_path = os.path.join(pdf_dev_path, 'mineru')
        local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
        image_dir = str(os.path.basename(local_image_dir))
        image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
        common.delete_file(dir_path)
        dss = read_local_images(pdf_path, suffixes=['.png', '.jpg'])
        count = 0
        for ds in dss:
            ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{count}.md", image_dir)
            count += 1
        common.sdk_count_folders_and_check_contents(dir_path)

    def test_local_doc_parse(self):
        """
        doc 解析
        """
        demo_names = list()
        pdf_path = os.path.join(pdf_dev_path, 'doc')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.docx'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            pdf_path = os.path.join(pdf_dev_path, 'doc', f'{demo_name}.docx')
            local_image_dir = os.path.join(pdf_dev_path, 'mineru', 'images')
            image_dir = str(os.path.basename(local_image_dir))
            name_without_suff = os.path.basename(pdf_path).split(".docx")[0]
            dir_path = os.path.join(pdf_dev_path, 'mineru')
            image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(dir_path)
            ds = read_local_office(pdf_path)[0]
            common.delete_file(dir_path)
            
            ds.apply(doc_analyze, ocr=True).pipe_txt_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir)          
            common.sdk_count_folders_and_check_contents(dir_path)


yyy's avatar
yyy committed
235
236
237
    @pytest.mark.P0
    def test_pdf_cli_auto(self):
        """magic_pdf cli test auto."""
yyy's avatar
yyy committed
238
        time.sleep(2)
yyy's avatar
yyy committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252
        demo_names = []
        pdf_path = os.path.join(pdf_dev_path, 'pdf')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.pdf'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            res_path = os.path.join(pdf_dev_path, 'mineru')
            common.delete_file(res_path)
            cmd = 'magic-pdf -p %s -o %s -m %s' % (os.path.join(
                pdf_path, f'{demo_name}.pdf'), res_path, 'auto')
            logging.info(cmd)
            os.system(cmd)
            common.cli_count_folders_and_check_contents(
                os.path.join(res_path, demo_name, 'auto'))
dt-yy's avatar
dt-yy committed
253
  
yyy's avatar
yyy committed
254
    @pytest.mark.P0
yyy's avatar
yyy committed
255
    def test_pdf_cli_txt(self):
yyy's avatar
yyy committed
256
        """magic_pdf cli test txt."""
yyy's avatar
yyy committed
257
        time.sleep(2)
yyy's avatar
yyy committed
258
259
260
261
262
263
264
265
266
267
268
269
270
271
        demo_names = []
        pdf_path = os.path.join(pdf_dev_path, 'pdf')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.pdf'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            res_path = os.path.join(pdf_dev_path, 'mineru')
            common.delete_file(res_path)
            cmd = 'magic-pdf -p %s -o %s -m %s' % (os.path.join(
                pdf_path, f'{demo_name}.pdf'), res_path, 'txt')
            logging.info(cmd)
            os.system(cmd)
            common.cli_count_folders_and_check_contents(
                os.path.join(res_path, demo_name, 'txt'))
yyy's avatar
yyy committed
272
   
yyy's avatar
yyy committed
273
    @pytest.mark.P0
yyy's avatar
yyy committed
274
    def test_pdf_cli_ocr(self):
yyy's avatar
yyy committed
275
        """magic_pdf cli test ocr."""
yyy's avatar
yyy committed
276
        time.sleep(2)
yyy's avatar
yyy committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
        demo_names = []
        pdf_path = os.path.join(pdf_dev_path, 'pdf')
        for pdf_file in os.listdir(pdf_path):
            if pdf_file.endswith('.pdf'):
                demo_names.append(pdf_file.split('.')[0])
        for demo_name in demo_names:
            res_path = os.path.join(pdf_dev_path, 'mineru')
            common.delete_file(res_path)
            cmd = 'magic-pdf -p %s -o %s -m %s' % (os.path.join(
                pdf_path, f'{demo_name}.pdf'), res_path, 'ocr')
            logging.info(cmd)
            os.system(cmd)
            common.cli_count_folders_and_check_contents(
                os.path.join(res_path, demo_name, 'ocr'))
yyy's avatar
yyy committed
291
292
    
    @pytest.mark.skip(reason='out-of-date api')
quyuan's avatar
quyuan committed
293
294
295
    @pytest.mark.P1
    def test_pdf_dev_cli_local_jsonl_txt(self):
        """magic_pdf_dev cli local txt."""
yyy's avatar
yyy committed
296
        time.sleep(2)
quyuan's avatar
quyuan committed
297
        jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
quyuan's avatar
quyuan committed
298
        cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, "txt")
quyuan's avatar
quyuan committed
299
300
301
        logging.info(cmd)
        os.system(cmd)

yyy's avatar
yyy committed
302
    @pytest.mark.skip(reason='out-of-date api')
quyuan's avatar
quyuan committed
303
304
305
    @pytest.mark.P1
    def test_pdf_dev_cli_local_jsonl_ocr(self):
        """magic_pdf_dev cli local ocr."""
yyy's avatar
yyy committed
306
        time.sleep(2)
quyuan's avatar
quyuan committed
307
        jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
quyuan's avatar
quyuan committed
308
        cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'ocr')
quyuan's avatar
quyuan committed
309
310
311
        logging.info(cmd)
        os.system(cmd)

yyy's avatar
yyy committed
312
    @pytest.mark.skip(reason='out-of-date api')
quyuan's avatar
quyuan committed
313
314
315
    @pytest.mark.P1
    def test_pdf_dev_cli_local_jsonl_auto(self):
        """magic_pdf_dev cli local auto."""
yyy's avatar
yyy committed
316
        time.sleep(2)
quyuan's avatar
quyuan committed
317
        jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
quyuan's avatar
quyuan committed
318
        cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'auto')
quyuan's avatar
quyuan committed
319
320
        logging.info(cmd)
        os.system(cmd)
yyy's avatar
yyy committed
321
322
    
    @pytest.mark.skip(reason='out-of-date api')
quyuan's avatar
quyuan committed
323
324
325
    @pytest.mark.P1
    def test_pdf_dev_cli_s3_jsonl_txt(self):
        """magic_pdf_dev cli s3 txt."""
yyy's avatar
yyy committed
326
        time.sleep(2)
quyuan's avatar
quyuan committed
327
        jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
quyuan's avatar
quyuan committed
328
        cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, "txt")
quyuan's avatar
quyuan committed
329
330
331
        logging.info(cmd)
        os.system(cmd)

yyy's avatar
yyy committed
332
    @pytest.mark.skip(reason='out-of-date api')
quyuan's avatar
quyuan committed
333
334
335
    @pytest.mark.P1
    def test_pdf_dev_cli_s3_jsonl_ocr(self):
        """magic_pdf_dev cli s3 ocr."""
yyy's avatar
yyy committed
336
        time.sleep(2)
quyuan's avatar
quyuan committed
337
        jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
quyuan's avatar
quyuan committed
338
        cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'ocr')
quyuan's avatar
quyuan committed
339
340
341
        logging.info(cmd)
        os.system(cmd)

yyy's avatar
yyy committed
342
    @pytest.mark.skip(reason='out-of-date api')
quyuan's avatar
quyuan committed
343
344
345
    @pytest.mark.P1
    def test_pdf_dev_cli_s3_jsonl_auto(self):
        """magic_pdf_dev cli s3 auto."""
yyy's avatar
yyy committed
346
        time.sleep(2)
quyuan's avatar
quyuan committed
347
348
349
350
351
        jsonl_path = os.path.join(pdf_dev_path, 'line1.jsonl')
        cmd = 'magic-pdf-dev --jsonl %s --method %s' % (jsonl_path, 'auto')
        logging.info(cmd)
        os.system(cmd)

quyuan's avatar
quyuan committed
352
353
354
    @pytest.mark.P1
    def test_pdf_dev_cli_pdf_json_auto(self):
        """magic_pdf_dev cli pdf+json auto."""
yyy's avatar
yyy committed
355
        time.sleep(2)
quyuan's avatar
quyuan committed
356
        json_path = os.path.join(pdf_dev_path, 'test_model.json')
yyy's avatar
yyy committed
357
        pdf_path = os.path.join(pdf_dev_path, 'pdf', 'test_rearch_report.pdf')
quyuan's avatar
quyuan committed
358
359
360
        cmd = 'magic-pdf-dev --pdf %s --json %s --method %s' % (pdf_path, json_path, 'auto')
        logging.info(cmd)
        os.system(cmd)
yyy's avatar
yyy committed
361
362
   
    @pytest.mark.skip(reason='out-of-date api')
quyuan's avatar
quyuan committed
363
364
365
    @pytest.mark.P1
    def test_pdf_dev_cli_pdf_json_ocr(self):
        """magic_pdf_dev cli pdf+json ocr."""
yyy's avatar
yyy committed
366
        time.sleep(2)
quyuan's avatar
quyuan committed
367
        json_path = os.path.join(pdf_dev_path, 'test_model.json')
yyy's avatar
yyy committed
368
        pdf_path = os.path.join(pdf_dev_path, 'pdf', 'test_rearch_report.pdf')
quyuan's avatar
quyuan committed
369
370
371
        cmd = 'magic-pdf-dev --pdf %s --json %s --method %s' % (pdf_path, json_path, 'auto')
        logging.info(cmd)
        os.system(cmd)
dt-yy's avatar
dt-yy committed
372
373
   

quyuan's avatar
quyuan committed
374
    @pytest.mark.P1
dt-yy's avatar
dt-yy committed
375
    def test_s3_sdk_auto(self):
yyy's avatar
yyy committed
376
377
378
379
380
381
        """
        test s3 sdk auto.
        """
        time.sleep(2)
        pdf_ak = os.getenv('pdf_ak')
        print (pdf_ak)
quyuan's avatar
quyuan committed
382
383
384
385
        pdf_sk = os.environ.get('pdf_sk', "")
        pdf_bucket = os.environ.get('bucket', "")
        pdf_endpoint = os.environ.get('pdf_endpoint', "")
        s3_pdf_path = conf.conf["s3_pdf_path"]
yyy's avatar
yyy committed
386
        image_dir = "s3://" + pdf_bucket + "/mineru/test/output"
dt-yy's avatar
dt-yy committed
387
388
        prefix = "mineru/test/output"
        reader = S3DataReader(prefix, pdf_bucket, pdf_ak, pdf_sk, pdf_endpoint)
dt-yy's avatar
dt-yy committed
389
        writer = S3DataWriter(prefix, pdf_bucket, pdf_ak, pdf_sk, pdf_endpoint)
dt-yy's avatar
dt-yy committed
390
391
        # = S3DataWriter(prefix, pdf_bucket, pdf_ak, pdf_sk, pdf_endpoint)
        image_writer = S3DataWriter(prefix, pdf_bucket, pdf_ak, pdf_sk, pdf_endpoint)
dt-yy's avatar
dt-yy committed
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
        local_dir = "output"
        name_without_suff = os.path.basename(s3_pdf_path).split(".")[0]

        # read bytes
        pdf_bytes = reader.read(s3_pdf_path)  # read the pdf content

        # proc
        ## Create Dataset Instance
        ds = PymuDocDataset(pdf_bytes)

        ## inference
        if ds.classify() == SupportedPdfParseMethod.OCR:
            infer_result = ds.apply(doc_analyze, ocr=True)

            ## pipeline
            pipe_result = infer_result.pipe_ocr_mode(image_writer)
        else:
            infer_result = ds.apply(doc_analyze, ocr=False)

            ## pipeline
            pipe_result = infer_result.pipe_txt_mode(image_writer)

        ### draw model result on each page
        infer_result.draw_model(os.path.join(local_dir, f'{name_without_suff}_model.pdf'))  # dump to local

        ### 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

        ### dump markdown
        pipe_result.dump_md(writer, f'{name_without_suff}.md', "unittest/tmp/images")    # dump to remote s3

        ### dump content list
        pipe_result.dump_content_list(writer, f"{name_without_suff}_content_list.json", image_dir)

quyuan's avatar
quyuan committed
429

yyy's avatar
yyy committed
430
431
432
433
    @pytest.mark.P1
    def test_local_magic_pdf_open_st_table(self):
        """magic pdf cli open st table."""
        time.sleep(2)
quyuan's avatar
quyuan committed
434
435
436
437
438
439
440
        #pre_cmd = "cp ~/magic_pdf_st.json ~/magic-pdf.json"
        value = {
        "model": "struct_eqtable",
        "enable": True,
        "max_time": 400
        }   
        common.update_config_file(magic_pdf_config, "table-config", value)
yyy's avatar
yyy committed
441
442
443
444
        pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
        common.delete_file(pdf_res_path)
        cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
        os.system(cli_cmd)
quyuan's avatar
quyuan committed
445
        res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
yyy's avatar
yyy committed
446
447
448
        assert res is True
  
    @pytest.mark.P1
quyuan's avatar
quyuan committed
449
450
    def test_local_magic_pdf_open_tablemaster_cuda(self):
        """magic pdf cli open table master html table cuda mode."""
yyy's avatar
yyy committed
451
        time.sleep(2)
quyuan's avatar
quyuan committed
452
453
454
455
456
457
458
459
        #pre_cmd = "cp ~/magic_pdf_html.json ~/magic-pdf.json"
        #os.system(pre_cmd)
        value = {
        "model": "tablemaster",
        "enable": True,
        "max_time": 400
        }   
        common.update_config_file(magic_pdf_config, "table-config", value)
yyy's avatar
yyy committed
460
461
462
463
464
465
466
467
        pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
        common.delete_file(pdf_res_path)
        cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
        os.system(cli_cmd)
        res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
        assert res is True
    
    @pytest.mark.P1
quyuan's avatar
quyuan committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
    def test_local_magic_pdf_open_rapidai_table(self):
        """magic pdf cli open rapid ai table."""
        time.sleep(2)
        #pre_cmd = "cp ~/magic_pdf_html.json ~/magic-pdf.json"
        #os.system(pre_cmd)
        value = {
        "model": "rapid_table",
        "enable": True,
        "max_time": 400
        }   
        common.update_config_file(magic_pdf_config, "table-config", value)
        pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
        common.delete_file(pdf_res_path)
        cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
        os.system(cli_cmd)
        res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
        assert res is True
    
    
    @pytest.mark.P1
    def test_local_magic_pdf_doclayout_yolo(self):
        """magic pdf cli open doclyaout yolo."""
        time.sleep(2)
        #pre_cmd = "cp ~/magic_pdf_html.json ~/magic-pdf.json"
        #os.system(pre_cmd)
        value = {
        "model": "doclayout_yolo"
        }   
        common.update_config_file(magic_pdf_config, "layout-config", value)
        pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
        common.delete_file(pdf_res_path)
        cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
        os.system(cli_cmd)
        common.cli_count_folders_and_check_contents(os.path.join(pdf_res_path, "test_rearch_report", "auto"))

    @pytest.mark.P1
    def test_local_magic_pdf_layoutlmv3_yolo(self):
        """magic pdf cli open layoutlmv3."""
        time.sleep(2)
        value = {
        "model": "layoutlmv3"
        }   
        common.update_config_file(magic_pdf_config, "layout-config", value)
        pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
        common.delete_file(pdf_res_path)
        cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
        os.system(cli_cmd)
        common.cli_count_folders_and_check_contents(os.path.join(pdf_res_path, "test_rearch_report", "auto"))
        #res = common.check_html_table_exists(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))

    @pytest.mark.P1
quyuan's avatar
quyuan committed
519
520
    def test_magic_pdf_cpu(self):
        """magic pdf cli cpu mode."""
yyy's avatar
yyy committed
521
        time.sleep(2)
quyuan's avatar
quyuan committed
522
523
524
525
        #pre_cmd = "cp ~/magic_pdf_html_table_cpu.json ~/magic-pdf.json"
        #os.system(pre_cmd)
        value = {
        "model": "tablemaster",
quyuan's avatar
quyuan committed
526
        "enable": False,
quyuan's avatar
quyuan committed
527
528
529
530
        "max_time": 400
        }   
        common.update_config_file(magic_pdf_config, "table-config", value)
        common.update_config_file(magic_pdf_config, "device-mode", "cpu")
yyy's avatar
yyy committed
531
532
533
534
        pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
        common.delete_file(pdf_res_path)
        cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
        os.system(cli_cmd)
quyuan's avatar
quyuan committed
535
        common.cli_count_folders_and_check_contents(os.path.join(pdf_res_path, "test_rearch_report", "auto"))
yyy's avatar
yyy committed
536

quyuan's avatar
quyuan committed
537

yyy's avatar
yyy committed
538
539
540
541
    @pytest.mark.P1
    def test_local_magic_pdf_close_html_table(self):
        """magic pdf cli close table."""
        time.sleep(2)
quyuan's avatar
quyuan committed
542
543
544
545
546
547
548
549
        #pre_cmd = "cp ~/magic_pdf_close_table.json ~/magic-pdf.json"
        #os.system(pre_cmd)
        value = {
        "model": "tablemaster",
        "enable": False,
        "max_time": 400
        }   
        common.update_config_file(magic_pdf_config, "table-config", value)
yyy's avatar
yyy committed
550
551
552
553
554
555
556
        pdf_path = os.path.join(pdf_dev_path, "pdf", "test_rearch_report.pdf")
        common.delete_file(pdf_res_path)
        cli_cmd = "magic-pdf -p %s -o %s" % (pdf_path, pdf_res_path)
        os.system(cli_cmd)
        res = common.check_close_tables(os.path.join(pdf_res_path, "test_rearch_report", "auto", "test_rearch_report.md"))
        assert res is True
    
quyuan's avatar
quyuan committed
557

yyy's avatar
yyy committed
558
 
yyy's avatar
yyy committed
559
560
if __name__ == '__main__':
    pytest.main()
dt-yy's avatar
dt-yy committed
561