UNIPipe.py 2.6 KB
Newer Older
赵小蒙's avatar
赵小蒙 committed
1
2
import json

3
from loguru import logger
kernel.h@qq.com's avatar
kernel.h@qq.com committed
4
5
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
赵小蒙's avatar
赵小蒙 committed
6
from magic_pdf.libs.commons import join_path
赵小蒙's avatar
赵小蒙 committed
7
from magic_pdf.pipe.AbsPipe import AbsPipe
赵小蒙's avatar
赵小蒙 committed
8
from magic_pdf.user_api import parse_union_pdf, parse_ocr_pdf
9
10


赵小蒙's avatar
赵小蒙 committed
11
class UNIPipe(AbsPipe):
12

赵小蒙's avatar
赵小蒙 committed
13
14
    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, img_parent_path: str,
                 is_debug: bool = False):
kernel.h@qq.com's avatar
kernel.h@qq.com committed
15
        self.pdf_type = self.PIP_OCR
kernel.h@qq.com's avatar
kernel.h@qq.com committed
16
        super().__init__(pdf_bytes, model_list, image_writer, img_parent_path, is_debug)
17

赵小蒙's avatar
赵小蒙 committed
18
19
    def pipe_classify(self):
        self.pdf_type = UNIPipe.classify(self.pdf_bytes)
20

赵小蒙's avatar
赵小蒙 committed
21
    def pipe_parse(self):
kernel.h@qq.com's avatar
kernel.h@qq.com committed
22
        if self.pdf_type == self.PIP_TXT:
赵小蒙's avatar
赵小蒙 committed
23
24
            self.pdf_mid_data = parse_union_pdf(self.pdf_bytes, self.model_list, self.image_writer,
                                                is_debug=self.is_debug)
kernel.h@qq.com's avatar
kernel.h@qq.com committed
25
        elif self.pdf_type == self.PIP_OCR:
赵小蒙's avatar
赵小蒙 committed
26
27
            self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer,
                                              is_debug=self.is_debug)
28

赵小蒙's avatar
赵小蒙 committed
29
    def pipe_mk_uni_format(self):
kernel.h@qq.com's avatar
kernel.h@qq.com committed
30
        content_list = AbsPipe.mk_uni_format(self.get_compress_pdf_mid_data(), self.img_parent_path)
赵小蒙's avatar
赵小蒙 committed
31
        return content_list
赵小蒙's avatar
赵小蒙 committed
32

赵小蒙's avatar
赵小蒙 committed
33
    def pipe_mk_markdown(self):
kernel.h@qq.com's avatar
kernel.h@qq.com committed
34
        markdown_content = AbsPipe.mk_markdown(self.get_compress_pdf_mid_data(), self.img_parent_path)
赵小蒙's avatar
赵小蒙 committed
35
        return markdown_content
36

赵小蒙's avatar
赵小蒙 committed
37

38
39
if __name__ == '__main__':
    # 测试
赵小蒙's avatar
赵小蒙 committed
40
41
    drw = DiskReaderWriter(r"D:/project/20231108code-clean")

赵小蒙's avatar
赵小蒙 committed
42
43
    pdf_file_path = r"linshixuqiu\19983-00.pdf"
    model_file_path = r"linshixuqiu\19983-00.json"
赵小蒙's avatar
赵小蒙 committed
44
45
    pdf_bytes = drw.read(pdf_file_path, AbsReaderWriter.MODE_BIN)
    model_json_txt = drw.read(model_file_path, AbsReaderWriter.MODE_TXT)
赵小蒙's avatar
赵小蒙 committed
46
47
48
49
    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))
赵小蒙's avatar
赵小蒙 committed
50

赵小蒙's avatar
赵小蒙 committed
51
52
53
54
55
56
57
58
    pipe = UNIPipe(pdf_bytes, model_list, img_writer, img_bucket_path)
    pipe.pipe_classify()
    pipe.pipe_parse()
    md_content = pipe.pipe_mk_markdown()
    try:
        content_list = pipe.pipe_mk_uni_format()
    except Exception as e:
        logger.exception(e)
赵小蒙's avatar
赵小蒙 committed
59
60

    md_writer = DiskReaderWriter(write_path)
赵小蒙's avatar
赵小蒙 committed
61
    md_writer.write(md_content, "19983-00.md", AbsReaderWriter.MODE_TXT)
赵小蒙's avatar
赵小蒙 committed
62
63
    md_writer.write(json.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4), "19983-00.json",
                    AbsReaderWriter.MODE_TXT)
赵小蒙's avatar
赵小蒙 committed
64
    md_writer.write(str(content_list), "19983-00.txt", AbsReaderWriter.MODE_TXT)