diff --git a/magic_pdf/pipe/AbsPipe.py b/magic_pdf/pipe/AbsPipe.py index 3580acd7..eb0b74cf 100644 --- a/magic_pdf/pipe/AbsPipe.py +++ b/magic_pdf/pipe/AbsPipe.py @@ -78,9 +78,9 @@ class AbsPipe(ABC): pdf_meta["text_layout_per_page"], ) if is_text_pdf: - return "txt" + return AbsPipe.PIP_TXT else: - return "ocr" + return AbsPipe.PIP_OCR @staticmethod def mk_uni_format(compressed_pdf_mid_data: str, img_buket_path: str) -> list: @@ -90,9 +90,9 @@ class AbsPipe(ABC): pdf_mid_data = JsonCompressor.decompress_json(compressed_pdf_mid_data) parse_type = pdf_mid_data["_parse_type"] pdf_info_list = pdf_mid_data["pdf_info"] - if parse_type == "txt": + if parse_type == AbsPipe.PIP_TXT: content_list = mk_universal_format(pdf_info_list, img_buket_path) - elif parse_type == "ocr": + elif parse_type == AbsPipe.PIP_OCR: content_list = make_standard_format_with_para(pdf_info_list, img_buket_path) return content_list @@ -104,10 +104,10 @@ class AbsPipe(ABC): pdf_mid_data = JsonCompressor.decompress_json(compressed_pdf_mid_data) parse_type = pdf_mid_data["_parse_type"] pdf_info_list = pdf_mid_data["pdf_info"] - if parse_type == "txt": + if parse_type == AbsPipe.PIP_TXT: content_list = mk_universal_format(pdf_info_list, img_buket_path) md_content = mk_mm_markdown(content_list) - elif parse_type == "ocr": + elif parse_type == AbsPipe.PIP_OCR: md_content = ocr_mk_mm_markdown_with_para(pdf_info_list, img_buket_path) return md_content diff --git a/magic_pdf/user_api.py b/magic_pdf/user_api.py index 71602f11..04ece100 100644 --- a/magic_pdf/user_api.py +++ b/magic_pdf/user_api.py @@ -19,6 +19,9 @@ from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr from magic_pdf.pdf_parse_by_txt import parse_pdf_by_txt +PARSE_TYPE_TXT = "txt" +PARSE_TYPE_OCR = "ocr" + def parse_txt_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: AbsReaderWriter, is_debug=False, start_page=0, *args, **kwargs): """ @@ -32,7 +35,7 @@ def parse_txt_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: AbsReaderWrit debug_mode=is_debug, ) - pdf_info_dict["parse_type"] = "txt" + pdf_info_dict["_parse_type"] = PARSE_TYPE_TXT return pdf_info_dict @@ -50,7 +53,7 @@ def parse_ocr_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: AbsReaderWrit debug_mode=is_debug, ) - pdf_info_dict["_parse_type"] = "ocr" + pdf_info_dict["_parse_type"] = PARSE_TYPE_OCR return pdf_info_dict @@ -82,8 +85,8 @@ def parse_union_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: AbsReaderWr if pdf_info_dict is None: raise Exception("Both parse_pdf_by_txt and parse_pdf_by_ocr failed.") else: - pdf_info_dict["_parse_type"] = "ocr" + pdf_info_dict["_parse_type"] = PARSE_TYPE_OCR else: - pdf_info_dict["_parse_type"] = "txt" + pdf_info_dict["_parse_type"] = PARSE_TYPE_TXT return pdf_info_dict