From b1fe9d4f60c03b436551e55ba0b575519c592964 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 3 Dec 2024 11:44:25 +0800 Subject: [PATCH 01/50] feat(gradio_app): implement dynamic concurrency limit based on VRAM - Add get_concurrency_limit function to calculate concurrency limit based on VRAM - Update clean_vram function and rename to get_vram for better clarity - Apply concurrency limit to the to_markdown function in the Gradio app --- magic_pdf/model/sub_modules/model_utils.py | 16 +++++++++++----- projects/gradio_app/app.py | 13 ++++++++++++- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/magic_pdf/model/sub_modules/model_utils.py b/magic_pdf/model/sub_modules/model_utils.py index 55114679..c57f50fa 100644 --- a/magic_pdf/model/sub_modules/model_utils.py +++ b/magic_pdf/model/sub_modules/model_utils.py @@ -42,10 +42,16 @@ def get_res_list_from_layout_res(layout_res): def clean_vram(device, vram_threshold=8): + total_memory = get_vram(device) + if total_memory <= vram_threshold: + gc_start = time.time() + clean_memory() + gc_time = round(time.time() - gc_start, 2) + logger.info(f"gc time: {gc_time}") + + +def get_vram(device): if torch.cuda.is_available() and device != 'cpu': total_memory = torch.cuda.get_device_properties(device).total_memory / (1024 ** 3) # 将字节转换为 GB - if total_memory <= vram_threshold: - gc_start = time.time() - clean_memory() - gc_time = round(time.time() - gc_start, 2) - logger.info(f"gc time: {gc_time}") \ No newline at end of file + return total_memory + return 0 \ No newline at end of file diff --git a/projects/gradio_app/app.py b/projects/gradio_app/app.py index 1ccede3b..b4436021 100644 --- a/projects/gradio_app/app.py +++ b/projects/gradio_app/app.py @@ -14,7 +14,9 @@ from gradio_pdf import PDF from loguru import logger from magic_pdf.data.data_reader_writer import FileBasedDataReader +from magic_pdf.libs.config_reader import get_device from magic_pdf.libs.hash_utils import compute_sha256 +from magic_pdf.model.sub_modules.model_utils import get_vram from magic_pdf.tools.common import do_parse, prepare_env @@ -183,6 +185,15 @@ def to_pdf(file_path): return tmp_file_path +def get_concurrency_limit(vram_threshold=7.5): + vram = get_vram(device = get_device()) + concurrency_limit = int(vram // vram_threshold) + if concurrency_limit < 1: + concurrency_limit = 1 + # logger.info(f'concurrency_limit: {concurrency_limit}') + return concurrency_limit + + if __name__ == '__main__': with gr.Blocks() as demo: gr.HTML(header) @@ -219,7 +230,7 @@ if __name__ == '__main__': md_text = gr.TextArea(lines=45, show_copy_button=True) file.upload(fn=to_pdf, inputs=file, outputs=pdf_show) change_bu.click(fn=to_markdown, inputs=[pdf_show, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language], - outputs=[md, md_text, output_file, pdf_show]) + outputs=[md, md_text, output_file, pdf_show], concurrency_limit=get_concurrency_limit()) clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr, table_enable, language]) demo.launch(server_name='0.0.0.0') From a3a720ea87a5200d8f2e2f0500e5f1f50f613527 Mon Sep 17 00:00:00 2001 From: icecraft Date: Wed, 27 Nov 2024 14:49:41 +0800 Subject: [PATCH 02/50] refactor: isolate inference and pipeline --- magic_pdf/data/dataset.py | 69 +++++++- magic_pdf/filter/__init__.py | 32 ++++ magic_pdf/libs/draw_bbox.py | 20 ++- .../model/doc_analyze_by_custom_model.py | 166 ++++++++++------- magic_pdf/model/types.py | 122 +++++++++++++ magic_pdf/pdf_parse_by_ocr.py | 9 +- magic_pdf/pdf_parse_by_txt.py | 9 +- magic_pdf/pdf_parse_union_core_v2.py | 2 +- magic_pdf/pipe/AbsPipe.py | 5 +- magic_pdf/pipe/OCRPipe.py | 69 ++++++-- magic_pdf/pipe/TXTPipe.py | 9 +- magic_pdf/pipe/UNIPipe.py | 110 +++++++++--- magic_pdf/pipe/types.py | 62 +++++++ magic_pdf/tools/common.py | 167 +++++++++++------- magic_pdf/user_api.py | 63 +++++-- 15 files changed, 701 insertions(+), 213 deletions(-) create mode 100644 magic_pdf/model/types.py create mode 100644 magic_pdf/pipe/types.py diff --git a/magic_pdf/data/dataset.py b/magic_pdf/data/dataset.py index 0eee3c68..d71d6b84 100644 --- a/magic_pdf/data/dataset.py +++ b/magic_pdf/data/dataset.py @@ -1,11 +1,13 @@ +import os from abc import ABC, abstractmethod -from typing import Iterator +from typing import Callable, Iterator import fitz from magic_pdf.config.enums import SupportedPdfParseMethod from magic_pdf.data.schemas import PageInfo from magic_pdf.data.utils import fitz_doc_to_image +from magic_pdf.filter import classify class PageableData(ABC): @@ -28,6 +30,14 @@ class PageableData(ABC): """ pass + @abstractmethod + def draw_rect(self, rect_coords, color, fill, fill_opacity, width, overlay): + pass + + @abstractmethod + def insert_text(self, coord, content, fontsize, color): + pass + class Dataset(ABC): @abstractmethod @@ -66,6 +76,18 @@ class Dataset(ABC): """ pass + @abstractmethod + def dump_to_file(self, file_path: str): + pass + + @abstractmethod + def apply(self, proc: Callable, *args, **kwargs): + pass + + @abstractmethod + def classify(self) -> SupportedPdfParseMethod: + pass + class PymuDocDataset(Dataset): def __init__(self, bits: bytes): @@ -74,7 +96,8 @@ class PymuDocDataset(Dataset): Args: bits (bytes): the bytes of the pdf """ - self._records = [Doc(v) for v in fitz.open('pdf', bits)] + self._raw_fitz = fitz.open('pdf', bits) + self._records = [Doc(v) for v in self._raw_fitz] self._data_bits = bits self._raw_data = bits @@ -109,6 +132,19 @@ class PymuDocDataset(Dataset): """ return self._records[page_id] + def dump_to_file(self, file_path: str): + dir_name = os.path.dirname(file_path) + if dir_name not in ('', '.', '..'): + os.makedirs(dir_name, exist_ok=True) + self._raw_fitz.save(file_path) + + def apply(self, proc: Callable, *args, **kwargs): + new_args = tuple([self] + list(args)) + return proc(*new_args, **kwargs) + + def classify(self) -> SupportedPdfParseMethod: + return classify(self._data_bits) + class ImageDataset(Dataset): def __init__(self, bits: bytes): @@ -118,7 +154,8 @@ class ImageDataset(Dataset): bits (bytes): the bytes of the photo which will be converted to pdf first. then converted to pymudoc. """ pdf_bytes = fitz.open(stream=bits).convert_to_pdf() - self._records = [Doc(v) for v in fitz.open('pdf', pdf_bytes)] + self._raw_fitz = fitz.open('pdf', pdf_bytes) + self._records = [Doc(v) for v in self._raw_fitz] self._raw_data = bits self._data_bits = pdf_bytes @@ -153,9 +190,22 @@ class ImageDataset(Dataset): """ return self._records[page_id] + def dump_to_file(self, file_path: str): + dir_name = os.path.dirname(file_path) + if dir_name not in ('', '.', '..'): + os.makedirs(dir_name, exist_ok=True) + self._raw_fitz.save(file_path) + + def apply(self, proc: Callable, *args, **kwargs): + return proc(self, *args, **kwargs) + + def classify(self) -> SupportedPdfParseMethod: + return SupportedPdfParseMethod.OCR + class Doc(PageableData): """Initialized with pymudoc object.""" + def __init__(self, doc: fitz.Page): self._doc = doc @@ -192,3 +242,16 @@ class Doc(PageableData): def __getattr__(self, name): if hasattr(self._doc, name): return getattr(self._doc, name) + + def draw_rect(self, rect_coords, color, fill, fill_opacity, width, overlay): + self._doc.draw_rect( + rect_coords, + color=color, + fill=fill, + fill_opacity=fill_opacity, + width=width, + overlay=overlay, + ) + + def insert_text(self, coord, content, fontsize, color): + self._doc.insert_text(coord, content, fontsize=fontsize, color=color) diff --git a/magic_pdf/filter/__init__.py b/magic_pdf/filter/__init__.py index e69de29b..07a79b88 100644 --- a/magic_pdf/filter/__init__.py +++ b/magic_pdf/filter/__init__.py @@ -0,0 +1,32 @@ + +from magic_pdf.config.drop_reason import DropReason +from magic_pdf.config.enums import SupportedPdfParseMethod +from magic_pdf.filter.pdf_classify_by_type import classify as do_classify +from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan + + +def classify(pdf_bytes: bytes) -> SupportedPdfParseMethod: + """根据pdf的元数据，判断是文本pdf，还是ocr pdf.""" + pdf_meta = pdf_meta_scan(pdf_bytes) + if pdf_meta.get('_need_drop', False): # 如果返回了需要丢弃的标志，则抛出异常 + raise Exception(f"pdf meta_scan need_drop,reason is {pdf_meta['_drop_reason']}") + else: + is_encrypted = pdf_meta['is_encrypted'] + is_needs_password = pdf_meta['is_needs_password'] + if is_encrypted or is_needs_password: # 加密的，需要密码的，没有页面的，都不处理 + raise Exception(f'pdf meta_scan need_drop,reason is {DropReason.ENCRYPTED}') + else: + is_text_pdf, results = do_classify( + pdf_meta['total_page'], + pdf_meta['page_width_pts'], + pdf_meta['page_height_pts'], + pdf_meta['image_info_per_page'], + pdf_meta['text_len_per_page'], + pdf_meta['imgs_per_page'], + pdf_meta['text_layout_per_page'], + pdf_meta['invalid_chars'], + ) + if is_text_pdf: + return SupportedPdfParseMethod.TXT + else: + return SupportedPdfParseMethod.OCR diff --git a/magic_pdf/libs/draw_bbox.py b/magic_pdf/libs/draw_bbox.py index 3aa2031c..e325ee45 100644 --- a/magic_pdf/libs/draw_bbox.py +++ b/magic_pdf/libs/draw_bbox.py @@ -1,7 +1,9 @@ import fitz from magic_pdf.config.constants import CROSS_PAGE -from magic_pdf.config.ocr_content_type import BlockType, CategoryId, ContentType -from magic_pdf.data.dataset import PymuDocDataset +from magic_pdf.config.ocr_content_type import (BlockType, CategoryId, + ContentType) +from magic_pdf.data.dataset import Dataset +from magic_pdf.libs.commons import fitz # PyMuPDF from magic_pdf.model.magic_model import MagicModel @@ -194,7 +196,7 @@ def draw_layout_bbox(pdf_info, pdf_bytes, out_path, filename): ) # Save the PDF - pdf_docs.save(f'{out_path}/{filename}_layout.pdf') + pdf_docs.save(f'{out_path}/{filename}') def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename): @@ -282,18 +284,17 @@ def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename): draw_bbox_without_number(i, dropped_list, page, [158, 158, 158], False) # Save the PDF - pdf_docs.save(f'{out_path}/{filename}_spans.pdf') + pdf_docs.save(f'{out_path}/{filename}') -def draw_model_bbox(model_list: list, pdf_bytes, out_path, filename): +def draw_model_bbox(model_list, dataset: Dataset, out_path, filename): dropped_bbox_list = [] tables_body_list, tables_caption_list, tables_footnote_list = [], [], [] imgs_body_list, imgs_caption_list, imgs_footnote_list = [], [], [] titles_list = [] texts_list = [] interequations_list = [] - pdf_docs = fitz.open('pdf', pdf_bytes) - magic_model = MagicModel(model_list, PymuDocDataset(pdf_bytes)) + magic_model = MagicModel(model_list, dataset) for i in range(len(model_list)): page_dropped_list = [] tables_body, tables_caption, tables_footnote = [], [], [] @@ -337,7 +338,8 @@ def draw_model_bbox(model_list: list, pdf_bytes, out_path, filename): dropped_bbox_list.append(page_dropped_list) imgs_footnote_list.append(imgs_footnote) - for i, page in enumerate(pdf_docs): + for i in range(len(dataset)): + page = dataset.get_page(i) draw_bbox_with_number( i, dropped_bbox_list, page, [158, 158, 158], True ) # color ! @@ -352,7 +354,7 @@ def draw_model_bbox(model_list: list, pdf_bytes, out_path, filename): draw_bbox_with_number(i, interequations_list, page, [0, 255, 0], True) # Save the PDF - pdf_docs.save(f'{out_path}/{filename}_model.pdf') + dataset.dump_to_file(f'{out_path}/{filename}') def draw_line_sort_bbox(pdf_info, pdf_bytes, out_path, filename): diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index a3536a3b..aabc7d2c 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -1,14 +1,19 @@ + import time import fitz import numpy as np from loguru import logger -from magic_pdf.libs.clean_memory import clean_memory -from magic_pdf.libs.config_reader import get_local_models_dir, get_device, get_table_recog_config, get_layout_config, \ - get_formula_config -from magic_pdf.model.model_list import MODEL import magic_pdf.model as model_config +from magic_pdf.data.dataset import Dataset +from magic_pdf.libs.clean_memory import clean_memory +from magic_pdf.libs.config_reader import (get_device, get_formula_config, + get_layout_config, + get_local_models_dir, + get_table_recog_config) +from magic_pdf.model.model_list import MODEL +from magic_pdf.model.types import InferenceResult def dict_compare(d1, d2): @@ -19,25 +24,31 @@ def remove_duplicates_dicts(lst): unique_dicts = [] for dict_item in lst: if not any( - dict_compare(dict_item, existing_dict) for existing_dict in unique_dicts + dict_compare(dict_item, existing_dict) for existing_dict in unique_dicts ): unique_dicts.append(dict_item) return unique_dicts -def load_images_from_pdf(pdf_bytes: bytes, dpi=200, start_page_id=0, end_page_id=None) -> list: +def load_images_from_pdf( + pdf_bytes: bytes, dpi=200, start_page_id=0, end_page_id=None +) -> list: try: from PIL import Image except ImportError: - logger.error("Pillow not installed, please install by pip.") + logger.error('Pillow not installed, please install by pip.') exit(1) images = [] - with fitz.open("pdf", pdf_bytes) as doc: + with fitz.open('pdf', pdf_bytes) as doc: pdf_page_num = doc.page_count - end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else pdf_page_num - 1 + end_page_id = ( + end_page_id + if end_page_id is not None and end_page_id >= 0 + else pdf_page_num - 1 + ) if end_page_id > pdf_page_num - 1: - logger.warning("end_page_id is out of range, use images length") + logger.warning('end_page_id is out of range, use images length') end_page_id = pdf_page_num - 1 for index in range(0, doc.page_count): @@ -50,11 +61,11 @@ def load_images_from_pdf(pdf_bytes: bytes, dpi=200, start_page_id=0, end_page_id if pm.width > 4500 or pm.height > 4500: pm = page.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False) - img = Image.frombytes("RGB", (pm.width, pm.height), pm.samples) + img = Image.frombytes('RGB', (pm.width, pm.height), pm.samples) img = np.array(img) - img_dict = {"img": img, "width": pm.width, "height": pm.height} + img_dict = {'img': img, 'width': pm.width, 'height': pm.height} else: - img_dict = {"img": [], "width": 0, "height": 0} + img_dict = {'img': [], 'width': 0, 'height': 0} images.append(img_dict) return images @@ -69,117 +80,150 @@ class ModelSingleton: cls._instance = super().__new__(cls) return cls._instance - def get_model(self, ocr: bool, show_log: bool, lang=None, layout_model=None, formula_enable=None, table_enable=None): + def get_model( + self, + ocr: bool, + show_log: bool, + lang=None, + layout_model=None, + formula_enable=None, + table_enable=None, + ): key = (ocr, show_log, lang, layout_model, formula_enable, table_enable) if key not in self._models: - self._models[key] = custom_model_init(ocr=ocr, show_log=show_log, lang=lang, layout_model=layout_model, - formula_enable=formula_enable, table_enable=table_enable) + self._models[key] = custom_model_init( + ocr=ocr, + show_log=show_log, + lang=lang, + layout_model=layout_model, + formula_enable=formula_enable, + table_enable=table_enable, + ) return self._models[key] -def custom_model_init(ocr: bool = False, show_log: bool = False, lang=None, - layout_model=None, formula_enable=None, table_enable=None): +def custom_model_init( + ocr: bool = False, + show_log: bool = False, + lang=None, + layout_model=None, + formula_enable=None, + table_enable=None, +): model = None - if model_config.__model_mode__ == "lite": - logger.warning("The Lite mode is provided for developers to conduct testing only, and the output quality is " - "not guaranteed to be reliable.") + if model_config.__model_mode__ == 'lite': + logger.warning( + 'The Lite mode is provided for developers to conduct testing only, and the output quality is ' + 'not guaranteed to be reliable.' + ) model = MODEL.Paddle - elif model_config.__model_mode__ == "full": + elif model_config.__model_mode__ == 'full': model = MODEL.PEK if model_config.__use_inside_model__: model_init_start = time.time() if model == MODEL.Paddle: from magic_pdf.model.pp_structure_v2 import CustomPaddleModel + custom_model = CustomPaddleModel(ocr=ocr, show_log=show_log, lang=lang) elif model == MODEL.PEK: from magic_pdf.model.pdf_extract_kit import CustomPEKModel + # 从配置文件读取model-dir和device local_models_dir = get_local_models_dir() device = get_device() layout_config = get_layout_config() if layout_model is not None: - layout_config["model"] = layout_model + layout_config['model'] = layout_model formula_config = get_formula_config() if formula_enable is not None: - formula_config["enable"] = formula_enable + formula_config['enable'] = formula_enable table_config = get_table_recog_config() if table_enable is not None: - table_config["enable"] = table_enable + table_config['enable'] = table_enable model_input = { - "ocr": ocr, - "show_log": show_log, - "models_dir": local_models_dir, - "device": device, - "table_config": table_config, - "layout_config": layout_config, - "formula_config": formula_config, - "lang": lang, + 'ocr': ocr, + 'show_log': show_log, + 'models_dir': local_models_dir, + 'device': device, + 'table_config': table_config, + 'layout_config': layout_config, + 'formula_config': formula_config, + 'lang': lang, } custom_model = CustomPEKModel(**model_input) else: - logger.error("Not allow model_name!") + logger.error('Not allow model_name!') exit(1) model_init_cost = time.time() - model_init_start - logger.info(f"model init cost: {model_init_cost}") + logger.info(f'model init cost: {model_init_cost}') else: - logger.error("use_inside_model is False, not allow to use inside model") + logger.error('use_inside_model is False, not allow to use inside model') exit(1) return custom_model -def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False, - start_page_id=0, end_page_id=None, lang=None, - layout_model=None, formula_enable=None, table_enable=None): +def doc_analyze( + dataset: Dataset, + ocr: bool = False, + show_log: bool = False, + start_page_id=0, + end_page_id=None, + lang=None, + layout_model=None, + formula_enable=None, + table_enable=None, +) -> InferenceResult: - if lang == "": + if lang == '': lang = None model_manager = ModelSingleton() - custom_model = model_manager.get_model(ocr, show_log, lang, layout_model, formula_enable, table_enable) - - with fitz.open("pdf", pdf_bytes) as doc: - pdf_page_num = doc.page_count - end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else pdf_page_num - 1 - if end_page_id > pdf_page_num - 1: - logger.warning("end_page_id is out of range, use images length") - end_page_id = pdf_page_num - 1 - - images = load_images_from_pdf(pdf_bytes, start_page_id=start_page_id, end_page_id=end_page_id) + custom_model = model_manager.get_model( + ocr, show_log, lang, layout_model, formula_enable, table_enable + ) model_json = [] doc_analyze_start = time.time() - for index, img_dict in enumerate(images): - img = img_dict["img"] - page_width = img_dict["width"] - page_height = img_dict["height"] + if end_page_id is None: + end_page_id = len(dataset) + + for index in range(len(dataset)): + page_data = dataset.get_page(index) + img_dict = page_data.get_image() + img = img_dict['img'] + page_width = img_dict['width'] + page_height = img_dict['height'] if start_page_id <= index <= end_page_id: page_start = time.time() result = custom_model(img) logger.info(f'-----page_id : {index}, page total time: {round(time.time() - page_start, 2)}-----') else: result = [] - page_info = {"page_no": index, "height": page_height, "width": page_width} - page_dict = {"layout_dets": result, "page_info": page_info} + + page_info = {'page_no': index, 'height': page_height, 'width': page_width} + page_dict = {'layout_dets': result, 'page_info': page_info} model_json.append(page_dict) gc_start = time.time() clean_memory() gc_time = round(time.time() - gc_start, 2) - logger.info(f"gc time: {gc_time}") + logger.info(f'gc time: {gc_time}') doc_analyze_time = round(time.time() - doc_analyze_start, 2) - doc_analyze_speed = round( (end_page_id + 1 - start_page_id) / doc_analyze_time, 2) - logger.info(f"doc analyze time: {round(time.time() - doc_analyze_start, 2)}," - f" speed: {doc_analyze_speed} pages/second") + doc_analyze_speed = round((end_page_id + 1 - start_page_id) / doc_analyze_time, 2) + logger.info( + f'doc analyze time: {round(time.time() - doc_analyze_start, 2)},' + f' speed: {doc_analyze_speed} pages/second' + ) - return model_json + return InferenceResult(model_json, dataset) diff --git a/magic_pdf/model/types.py b/magic_pdf/model/types.py new file mode 100644 index 00000000..4514d251 --- /dev/null +++ b/magic_pdf/model/types.py @@ -0,0 +1,122 @@ +import copy +import json +import os +from typing import Callable + +from magic_pdf.config.enums import SupportedPdfParseMethod +from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset +from magic_pdf.filter import classify +from magic_pdf.libs.draw_bbox import draw_model_bbox +from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union +from magic_pdf.pipe.types import PipeResult + + +class InferenceResult: + def __init__(self, inference_results: list, dataset: Dataset): + self._infer_res = inference_results + self._dataset = dataset + + def draw_model(self, file_path: str) -> None: + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + draw_model_bbox( + copy.deepcopy(self._infer_res), self._dataset, dir_name, base_name + ) + + def dump_model(self, writer: DataWriter, file_path: str): + writer.write_string( + file_path, json.dumps(self._infer_res, ensure_ascii=False, indent=4) + ) + + def get_infer_res(self): + return self._infer_res + + def apply(self, proc: Callable, *args, **kwargs): + return proc(copy.deepcopy(self._infer_res), *args, **kwargs) + + def pipe_auto_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + def proc(*args, **kwargs) -> PipeResult: + res = pdf_parse_union(*args, **kwargs) + return PipeResult(res, self._dataset) + + pdf_proc_method = classify(self._dataset.data_bits()) + + if pdf_proc_method == SupportedPdfParseMethod.TXT: + return self.apply( + proc, + self._dataset, + imageWriter, + SupportedPdfParseMethod.TXT, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) + else: + return self.apply( + proc, + self._dataset, + imageWriter, + SupportedPdfParseMethod.OCR, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) + + def pipe_txt_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + def proc(*args, **kwargs) -> PipeResult: + res = pdf_parse_union(*args, **kwargs) + return PipeResult(res, self._dataset) + + return self.apply( + proc, + self._dataset, + imageWriter, + SupportedPdfParseMethod.TXT, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) + + def pipe_ocr_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + + def proc(*args, **kwargs) -> PipeResult: + res = pdf_parse_union(*args, **kwargs) + return PipeResult(res, self._dataset) + + return self.apply( + proc, + self._dataset, + imageWriter, + SupportedPdfParseMethod.TXT, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) diff --git a/magic_pdf/pdf_parse_by_ocr.py b/magic_pdf/pdf_parse_by_ocr.py index 7f0cacc1..08228080 100644 --- a/magic_pdf/pdf_parse_by_ocr.py +++ b/magic_pdf/pdf_parse_by_ocr.py @@ -1,9 +1,9 @@ from magic_pdf.config.enums import SupportedPdfParseMethod -from magic_pdf.data.dataset import PymuDocDataset +from magic_pdf.data.dataset import Dataset from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union -def parse_pdf_by_ocr(pdf_bytes, +def parse_pdf_by_ocr(dataset: Dataset, model_list, imageWriter, start_page_id=0, @@ -11,9 +11,8 @@ def parse_pdf_by_ocr(pdf_bytes, debug_mode=False, lang=None, ): - dataset = PymuDocDataset(pdf_bytes) - return pdf_parse_union(dataset, - model_list, + return pdf_parse_union(model_list, + dataset, imageWriter, SupportedPdfParseMethod.OCR, start_page_id=start_page_id, diff --git a/magic_pdf/pdf_parse_by_txt.py b/magic_pdf/pdf_parse_by_txt.py index c7cd0787..5e896b67 100644 --- a/magic_pdf/pdf_parse_by_txt.py +++ b/magic_pdf/pdf_parse_by_txt.py @@ -1,10 +1,10 @@ from magic_pdf.config.enums import SupportedPdfParseMethod -from magic_pdf.data.dataset import PymuDocDataset +from magic_pdf.data.dataset import Dataset from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union def parse_pdf_by_txt( - pdf_bytes, + dataset: Dataset, model_list, imageWriter, start_page_id=0, @@ -12,9 +12,8 @@ def parse_pdf_by_txt( debug_mode=False, lang=None, ): - dataset = PymuDocDataset(pdf_bytes) - return pdf_parse_union(dataset, - model_list, + return pdf_parse_union(model_list, + dataset, imageWriter, SupportedPdfParseMethod.TXT, start_page_id=start_page_id, diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index 2e65191e..bd4b6c6a 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -832,4 +832,4 @@ def pdf_parse_union( if __name__ == '__main__': - pass + pass \ No newline at end of file diff --git a/magic_pdf/pipe/AbsPipe.py b/magic_pdf/pipe/AbsPipe.py index 91106a76..44d7212c 100644 --- a/magic_pdf/pipe/AbsPipe.py +++ b/magic_pdf/pipe/AbsPipe.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from magic_pdf.config.drop_reason import DropReason from magic_pdf.config.make_content_config import DropMode, MakeMode from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset from magic_pdf.dict2md.ocr_mkcontent import union_make from magic_pdf.filter.pdf_classify_by_type import classify from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan @@ -14,9 +15,9 @@ class AbsPipe(ABC): PIP_OCR = 'ocr' PIP_TXT = 'txt' - def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: DataWriter, is_debug: bool = False, + def __init__(self, dataset: Dataset, model_list: list, image_writer: DataWriter, is_debug: bool = False, start_page_id=0, end_page_id=None, lang=None, layout_model=None, formula_enable=None, table_enable=None): - self.pdf_bytes = pdf_bytes + self.dataset = Dataset self.model_list = model_list self.image_writer = image_writer self.pdf_mid_data = None # 未压缩 diff --git a/magic_pdf/pipe/OCRPipe.py b/magic_pdf/pipe/OCRPipe.py index 6368189c..3a52f5bd 100644 --- a/magic_pdf/pipe/OCRPipe.py +++ b/magic_pdf/pipe/OCRPipe.py @@ -2,40 +2,79 @@ from loguru import logger from magic_pdf.config.make_content_config import DropMode, MakeMode from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze from magic_pdf.pipe.AbsPipe import AbsPipe from magic_pdf.user_api import parse_ocr_pdf class OCRPipe(AbsPipe): - - def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: DataWriter, is_debug: bool = False, - start_page_id=0, end_page_id=None, lang=None, - layout_model=None, formula_enable=None, table_enable=None): - super().__init__(pdf_bytes, model_list, image_writer, is_debug, start_page_id, end_page_id, lang, - layout_model, formula_enable, table_enable) + def __init__( + self, + dataset: Dataset, + model_list: list, + image_writer: DataWriter, + is_debug: bool = False, + start_page_id=0, + end_page_id=None, + lang=None, + layout_model=None, + formula_enable=None, + table_enable=None, + ): + super().__init__( + dataset, + model_list, + image_writer, + is_debug, + start_page_id, + end_page_id, + lang, + layout_model, + formula_enable, + table_enable, + ) def pipe_classify(self): pass def pipe_analyze(self): - self.model_list = doc_analyze(self.pdf_bytes, ocr=True, - start_page_id=self.start_page_id, end_page_id=self.end_page_id, - lang=self.lang, layout_model=self.layout_model, - formula_enable=self.formula_enable, table_enable=self.table_enable) + self.infer_res = doc_analyze( + self.dataset, + ocr=True, + start_page_id=self.start_page_id, + end_page_id=self.end_page_id, + lang=self.lang, + layout_model=self.layout_model, + formula_enable=self.formula_enable, + table_enable=self.table_enable, + ) def pipe_parse(self): - self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer, is_debug=self.is_debug, - start_page_id=self.start_page_id, end_page_id=self.end_page_id, - lang=self.lang, layout_model=self.layout_model, - formula_enable=self.formula_enable, table_enable=self.table_enable) + self.pdf_mid_data = parse_ocr_pdf( + self.dataset, + self.infer_res, + self.image_writer, + is_debug=self.is_debug, + start_page_id=self.start_page_id, + end_page_id=self.end_page_id, + lang=self.lang, + layout_model=self.layout_model, + formula_enable=self.formula_enable, + table_enable=self.table_enable, + ) def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF): result = super().pipe_mk_uni_format(img_parent_path, drop_mode) logger.info('ocr_pipe mk content list finished') return result - def pipe_mk_markdown(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD): + def pipe_mk_markdown( + self, + img_parent_path: str, + drop_mode=DropMode.WHOLE_PDF, + md_make_mode=MakeMode.MM_MD, + ): result = super().pipe_mk_markdown(img_parent_path, drop_mode, md_make_mode) logger.info(f'ocr_pipe mk {md_make_mode} finished') return result diff --git a/magic_pdf/pipe/TXTPipe.py b/magic_pdf/pipe/TXTPipe.py index 50ce0842..99255b43 100644 --- a/magic_pdf/pipe/TXTPipe.py +++ b/magic_pdf/pipe/TXTPipe.py @@ -2,6 +2,7 @@ from loguru import logger from magic_pdf.config.make_content_config import DropMode, MakeMode from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze from magic_pdf.pipe.AbsPipe import AbsPipe from magic_pdf.user_api import parse_txt_pdf @@ -9,23 +10,23 @@ from magic_pdf.user_api import parse_txt_pdf class TXTPipe(AbsPipe): - def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: DataWriter, is_debug: bool = False, + def __init__(self, dataset: Dataset, model_list: list, image_writer: DataWriter, is_debug: bool = False, start_page_id=0, end_page_id=None, lang=None, layout_model=None, formula_enable=None, table_enable=None): - super().__init__(pdf_bytes, model_list, image_writer, is_debug, start_page_id, end_page_id, lang, + super().__init__(dataset, model_list, image_writer, is_debug, start_page_id, end_page_id, lang, layout_model, formula_enable, table_enable) def pipe_classify(self): pass def pipe_analyze(self): - self.model_list = doc_analyze(self.pdf_bytes, ocr=False, + self.model_list = doc_analyze(self.dataset, ocr=False, start_page_id=self.start_page_id, end_page_id=self.end_page_id, lang=self.lang, layout_model=self.layout_model, formula_enable=self.formula_enable, table_enable=self.table_enable) def pipe_parse(self): - self.pdf_mid_data = parse_txt_pdf(self.pdf_bytes, self.model_list, self.image_writer, is_debug=self.is_debug, + self.pdf_mid_data = parse_txt_pdf(self.dataset, self.model_list, self.image_writer, is_debug=self.is_debug, start_page_id=self.start_page_id, end_page_id=self.end_page_id, lang=self.lang, layout_model=self.layout_model, formula_enable=self.formula_enable, table_enable=self.table_enable) diff --git a/magic_pdf/pipe/UNIPipe.py b/magic_pdf/pipe/UNIPipe.py index ca498004..1809492e 100644 --- a/magic_pdf/pipe/UNIPipe.py +++ b/magic_pdf/pipe/UNIPipe.py @@ -4,6 +4,7 @@ from loguru import logger from magic_pdf.config.make_content_config import DropMode, MakeMode from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset from magic_pdf.libs.commons import join_path from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze from magic_pdf.pipe.AbsPipe import AbsPipe @@ -12,12 +13,32 @@ from magic_pdf.user_api import parse_ocr_pdf, parse_union_pdf class UNIPipe(AbsPipe): - def __init__(self, pdf_bytes: bytes, jso_useful_key: dict, image_writer: DataWriter, is_debug: bool = False, - start_page_id=0, end_page_id=None, lang=None, - layout_model=None, formula_enable=None, table_enable=None): + def __init__( + self, + dataset: Dataset, + jso_useful_key: dict, + image_writer: DataWriter, + is_debug: bool = False, + start_page_id=0, + end_page_id=None, + lang=None, + layout_model=None, + formula_enable=None, + table_enable=None, + ): self.pdf_type = jso_useful_key['_pdf_type'] - super().__init__(pdf_bytes, jso_useful_key['model_list'], image_writer, is_debug, start_page_id, end_page_id, - lang, layout_model, formula_enable, table_enable) + super().__init__( + dataset, + jso_useful_key['model_list'], + image_writer, + is_debug, + start_page_id, + end_page_id, + lang, + layout_model, + formula_enable, + table_enable, + ) if len(self.model_list) == 0: self.input_model_is_empty = True else: @@ -28,35 +49,66 @@ class UNIPipe(AbsPipe): def pipe_analyze(self): if self.pdf_type == self.PIP_TXT: - self.model_list = doc_analyze(self.pdf_bytes, ocr=False, - start_page_id=self.start_page_id, end_page_id=self.end_page_id, - lang=self.lang, layout_model=self.layout_model, - formula_enable=self.formula_enable, table_enable=self.table_enable) + self.model_list = doc_analyze( + self.dataset, + ocr=False, + start_page_id=self.start_page_id, + end_page_id=self.end_page_id, + lang=self.lang, + layout_model=self.layout_model, + formula_enable=self.formula_enable, + table_enable=self.table_enable, + ) elif self.pdf_type == self.PIP_OCR: - self.model_list = doc_analyze(self.pdf_bytes, ocr=True, - start_page_id=self.start_page_id, end_page_id=self.end_page_id, - lang=self.lang, layout_model=self.layout_model, - formula_enable=self.formula_enable, table_enable=self.table_enable) + self.model_list = doc_analyze( + self.dataset, + ocr=True, + start_page_id=self.start_page_id, + end_page_id=self.end_page_id, + lang=self.lang, + layout_model=self.layout_model, + formula_enable=self.formula_enable, + table_enable=self.table_enable, + ) def pipe_parse(self): if self.pdf_type == self.PIP_TXT: - self.pdf_mid_data = parse_union_pdf(self.pdf_bytes, self.model_list, self.image_writer, - is_debug=self.is_debug, input_model_is_empty=self.input_model_is_empty, - start_page_id=self.start_page_id, end_page_id=self.end_page_id, - lang=self.lang, layout_model=self.layout_model, - formula_enable=self.formula_enable, table_enable=self.table_enable) + self.pdf_mid_data = parse_union_pdf( + self.dataset, + self.model_list, + self.image_writer, + is_debug=self.is_debug, + start_page_id=self.start_page_id, + end_page_id=self.end_page_id, + lang=self.lang, + layout_model=self.layout_model, + formula_enable=self.formula_enable, + table_enable=self.table_enable, + ) elif self.pdf_type == self.PIP_OCR: - self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer, - is_debug=self.is_debug, - start_page_id=self.start_page_id, end_page_id=self.end_page_id, - lang=self.lang) + self.pdf_mid_data = parse_ocr_pdf( + self.dataset, + self.model_list, + self.image_writer, + is_debug=self.is_debug, + start_page_id=self.start_page_id, + end_page_id=self.end_page_id, + lang=self.lang, + ) - def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.NONE_WITH_REASON): + def pipe_mk_uni_format( + self, img_parent_path: str, drop_mode=DropMode.NONE_WITH_REASON + ): result = super().pipe_mk_uni_format(img_parent_path, drop_mode) logger.info('uni_pipe mk content list finished') return result - def pipe_mk_markdown(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD): + def pipe_mk_markdown( + self, + img_parent_path: str, + drop_mode=DropMode.WHOLE_PDF, + md_make_mode=MakeMode.MM_MD, + ): result = super().pipe_mk_markdown(img_parent_path, drop_mode, md_make_mode) logger.info(f'uni_pipe mk {md_make_mode} finished') return result @@ -65,6 +117,7 @@ class UNIPipe(AbsPipe): if __name__ == '__main__': # 测试 from magic_pdf.data.data_reader_writer import DataReader + drw = DataReader(r'D:/project/20231108code-clean') pdf_file_path = r'linshixuqiu\19983-00.pdf' @@ -82,10 +135,7 @@ if __name__ == '__main__': # "model_list": model_list # } - jso_useful_key = { - '_pdf_type': '', - 'model_list': model_list - } + jso_useful_key = {'_pdf_type': '', 'model_list': model_list} pipe = UNIPipe(pdf_bytes, jso_useful_key, img_writer) pipe.pipe_classify() pipe.pipe_parse() @@ -94,5 +144,7 @@ if __name__ == '__main__': md_writer = DataWriter(write_path) md_writer.write_string('19983-00.md', md_content) - md_writer.write_string('19983-00.json', json.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4)) + md_writer.write_string( + '19983-00.json', json.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4) + ) md_writer.write_string('19983-00.txt', str(content_list)) diff --git a/magic_pdf/pipe/types.py b/magic_pdf/pipe/types.py new file mode 100644 index 00000000..6ee9b2ce --- /dev/null +++ b/magic_pdf/pipe/types.py @@ -0,0 +1,62 @@ + +import json +import os + +from magic_pdf.config.make_content_config import DropMode, MakeMode +from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset +from magic_pdf.dict2md.ocr_mkcontent import union_make +from magic_pdf.libs.draw_bbox import (draw_layout_bbox, draw_line_sort_bbox, + draw_span_bbox) +from magic_pdf.libs.json_compressor import JsonCompressor + + +class PipeResult: + def __init__(self, pipe_res, dataset: Dataset): + self._pipe_res = pipe_res + self._dataset = dataset + + def dump_md(self, writer: DataWriter, file_path: str, img_dir_or_bucket_prefix: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD): + pdf_info_list = self._pipe_res['pdf_info'] + md_content = union_make(pdf_info_list, md_make_mode, drop_mode, img_dir_or_bucket_prefix) + writer.write_string(file_path, md_content) + + def dump_content_list(self, writer: DataWriter, file_path: str, image_dir_or_bucket_prefix: str, drop_mode=DropMode.NONE): + pdf_info_list = self._pipe_res['pdf_info'] + content_list = union_make(pdf_info_list, MakeMode.STANDARD_FORMAT, drop_mode, image_dir_or_bucket_prefix) + writer.write_string(file_path, json.dumps(content_list, ensure_ascii=False, indent=4)) + + def dump_middle_json(self, writer: DataWriter, file_path: str): + writer.write_string(file_path, json.dumps(self._pipe_res, ensure_ascii=False, indent=4)) + + def draw_layout(self, file_path: str) -> None: + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + pdf_info = self._pipe_res['pdf_info'] + draw_layout_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) + + def draw_span(self, file_path: str): + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + pdf_info = self._pipe_res['pdf_info'] + draw_span_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) + + def draw_line_sort(self, file_path: str): + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + pdf_info = self._pipe_res['pdf_info'] + draw_line_sort_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) + + def draw_content_list(self, writer: DataWriter, file_path: str, img_dir_or_bucket_prefix: str, drop_mode=DropMode.WHOLE_PDF): + pdf_info_list = self._pipe_res['pdf_info'] + content_list = union_make(pdf_info_list, MakeMode.STANDARD_FORMAT, drop_mode, img_dir_or_bucket_prefix) + writer.write_string(file_path, json.dumps(content_list, ensure_ascii=False, indent=4)) + + def get_compress_pdf_mid_data(self): + return JsonCompressor.compress_json(self.pdf_mid_data) diff --git a/magic_pdf/tools/common.py b/magic_pdf/tools/common.py index 8675d61c..3bd2443b 100644 --- a/magic_pdf/tools/common.py +++ b/magic_pdf/tools/common.py @@ -1,5 +1,3 @@ -import copy -import json as json_parse import os import click @@ -7,13 +5,12 @@ import fitz from loguru import logger import magic_pdf.model as model_config +from magic_pdf.config.enums import SupportedPdfParseMethod from magic_pdf.config.make_content_config import DropMode, MakeMode from magic_pdf.data.data_reader_writer import FileBasedDataWriter -from magic_pdf.libs.draw_bbox import (draw_layout_bbox, draw_line_sort_bbox, - draw_model_bbox, draw_span_bbox) -from magic_pdf.pipe.OCRPipe import OCRPipe -from magic_pdf.pipe.TXTPipe import TXTPipe -from magic_pdf.pipe.UNIPipe import UNIPipe +from magic_pdf.data.dataset import PymuDocDataset +from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze +from magic_pdf.model.types import InferenceResult # from io import BytesIO # from pypdf import PdfReader, PdfWriter @@ -56,7 +53,11 @@ def prepare_env(output_dir, pdf_file_name, method): def convert_pdf_bytes_to_bytes_by_pymupdf(pdf_bytes, start_page_id=0, end_page_id=None): document = fitz.open('pdf', pdf_bytes) output_document = fitz.open() - end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else len(document) - 1 + end_page_id = ( + end_page_id + if end_page_id is not None and end_page_id >= 0 + else len(document) - 1 + ) if end_page_id > len(document) - 1: logger.warning('end_page_id is out of range, use pdf_docs length') end_page_id = len(document) - 1 @@ -94,78 +95,123 @@ def do_parse( f_draw_model_bbox = True f_draw_line_sort_bbox = True - if lang == "": + if lang == '': lang = None - pdf_bytes = convert_pdf_bytes_to_bytes_by_pymupdf(pdf_bytes, start_page_id, end_page_id) + pdf_bytes = convert_pdf_bytes_to_bytes_by_pymupdf( + pdf_bytes, start_page_id, end_page_id + ) - orig_model_list = copy.deepcopy(model_list) - local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, - parse_method) + local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method) - image_writer, md_writer = FileBasedDataWriter( - local_image_dir), FileBasedDataWriter(local_md_dir) + image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter( + local_md_dir + ) image_dir = str(os.path.basename(local_image_dir)) - if parse_method == 'auto': - jso_useful_key = {'_pdf_type': '', 'model_list': model_list} - pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True, - # start_page_id=start_page_id, end_page_id=end_page_id, - lang=lang, - layout_model=layout_model, formula_enable=formula_enable, table_enable=table_enable) - elif parse_method == 'txt': - pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True, - # start_page_id=start_page_id, end_page_id=end_page_id, - lang=lang, - layout_model=layout_model, formula_enable=formula_enable, table_enable=table_enable) - elif parse_method == 'ocr': - pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True, - # start_page_id=start_page_id, end_page_id=end_page_id, - lang=lang, - layout_model=layout_model, formula_enable=formula_enable, table_enable=table_enable) - else: - logger.error('unknown parse method') - exit(1) - - pipe.pipe_classify() + ds = PymuDocDataset(pdf_bytes) if len(model_list) == 0: if model_config.__use_inside_model__: - pipe.pipe_analyze() - orig_model_list = copy.deepcopy(pipe.model_list) + if parse_method == 'auto': + if ds.classify() == SupportedPdfParseMethod.TXT: + infer_result = ds.apply( + doc_analyze, + ocr=False, + lang=lang, + layout_model=layout_model, + formula_enable=formula_enable, + table_enable=table_enable, + ) + else: + infer_result = ds.apply( + doc_analyze, + ocr=True, + lang=lang, + layout_model=layout_model, + formula_enable=formula_enable, + table_enable=table_enable, + ) + pipe_result = infer_result.pipe_auto_mode( + image_writer, debug_mode=True, lang=lang + ) + + elif parse_method == 'txt': + infer_result = ds.apply( + doc_analyze, + ocr=False, + lang=lang, + layout_model=layout_model, + formula_enable=formula_enable, + table_enable=table_enable, + ) + pipe_result = infer_result.pipe_txt_mode( + image_writer, debug_mode=True, lang=lang + ) + elif parse_method == 'ocr': + infer_result = ds.apply( + doc_analyze, + ocr=True, + lang=lang, + layout_model=layout_model, + formula_enable=formula_enable, + table_enable=table_enable, + ) + pipe_result = infer_result.pipe_ocr_mode( + image_writer, debug_mode=True, lang=lang + ) + else: + logger.error('unknown parse method') + exit(1) else: logger.error('need model list input') exit(2) + else: + infer_result = InferenceResult(model_list, ds) + if parse_method == 'ocr': + pipe_result = infer_result.pipe_ocr_mode( + image_writer, debug_mode=True, lang=lang + ) + elif parse_method == 'txt': + pipe_result = infer_result.pipe_txt_mode( + image_writer, debug_mode=True, lang=lang + ) + else: + pipe_result = infer_result.pipe_auto_mode( + image_writer, debug_mode=True, lang=lang + ) - pipe.pipe_parse() - pdf_info = pipe.pdf_mid_data['pdf_info'] - if f_draw_layout_bbox: - draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name) - if f_draw_span_bbox: - draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name) if f_draw_model_bbox: - draw_model_bbox(copy.deepcopy(orig_model_list), pdf_bytes, local_md_dir, pdf_file_name) + infer_result.draw_model( + os.path.join(local_md_dir, f'{pdf_file_name}_model.pdf') + ) + + if f_draw_layout_bbox: + pipe_result.draw_layout( + os.path.join(local_md_dir, f'{pdf_file_name}_layout.pdf') + ) + if f_draw_span_bbox: + pipe_result.draw_span(os.path.join(local_md_dir, f'{pdf_file_name}_spans.pdf')) + if f_draw_line_sort_bbox: - draw_line_sort_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name) + pipe_result.draw_line_sort( + os.path.join(local_md_dir, f'{pdf_file_name}_line_sort.pdf') + ) - md_content = pipe.pipe_mk_markdown(image_dir, drop_mode=DropMode.NONE, md_make_mode=f_make_md_mode) if f_dump_md: - md_writer.write_string( + pipe_result.dump_md( + md_writer, f'{pdf_file_name}.md', - md_content + image_dir, + drop_mode=DropMode.NONE, + md_make_mode=f_make_md_mode, ) if f_dump_middle_json: - md_writer.write_string( - f'{pdf_file_name}_middle.json', - json_parse.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4) - ) + pipe_result.dump_middle_json(md_writer, f'{pdf_file_name}_middle.json') if f_dump_model_json: - md_writer.write_string( - f'{pdf_file_name}_model.json', - json_parse.dumps(orig_model_list, ensure_ascii=False, indent=4) - ) + infer_result.dump_model(md_writer, f'{pdf_file_name}_model.json') if f_dump_orig_pdf: md_writer.write( @@ -173,11 +219,12 @@ def do_parse( pdf_bytes, ) - content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE) if f_dump_content_list: - md_writer.write_string( + pipe_result.dump_content_list( + md_writer, f'{pdf_file_name}_content_list.json', - json_parse.dumps(content_list, ensure_ascii=False, indent=4) + image_dir, + drop_mode=DropMode.NONE, ) logger.info(f'local output dir is {local_md_dir}') diff --git a/magic_pdf/user_api.py b/magic_pdf/user_api.py index 221900e3..99b9502f 100644 --- a/magic_pdf/user_api.py +++ b/magic_pdf/user_api.py @@ -10,6 +10,7 @@ from loguru import logger from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset from magic_pdf.libs.version import __version__ from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr @@ -19,13 +20,21 @@ PARSE_TYPE_TXT = 'txt' PARSE_TYPE_OCR = 'ocr' -def parse_txt_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: DataWriter, is_debug=False, - start_page_id=0, end_page_id=None, lang=None, - *args, **kwargs): +def parse_txt_pdf( + dataset: Dataset, + model_list: list, + imageWriter: DataWriter, + is_debug=False, + start_page_id=0, + end_page_id=None, + lang=None, + *args, + **kwargs +): """解析文本类pdf.""" pdf_info_dict = parse_pdf_by_txt( - pdf_bytes, - pdf_models, + dataset, + model_list, imageWriter, start_page_id=start_page_id, end_page_id=end_page_id, @@ -43,13 +52,21 @@ def parse_txt_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: DataWriter, i return pdf_info_dict -def parse_ocr_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: DataWriter, is_debug=False, - start_page_id=0, end_page_id=None, lang=None, - *args, **kwargs): +def parse_ocr_pdf( + dataset: Dataset, + model_list: list, + imageWriter: DataWriter, + is_debug=False, + start_page_id=0, + end_page_id=None, + lang=None, + *args, + **kwargs +): """解析ocr类pdf.""" pdf_info_dict = parse_pdf_by_ocr( - pdf_bytes, - pdf_models, + dataset, + model_list, imageWriter, start_page_id=start_page_id, end_page_id=end_page_id, @@ -67,17 +84,24 @@ def parse_ocr_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: DataWriter, i return pdf_info_dict -def parse_union_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: DataWriter, is_debug=False, - input_model_is_empty: bool = False, - start_page_id=0, end_page_id=None, lang=None, - *args, **kwargs): +def parse_union_pdf( + dataset: Dataset, + model_list: list, + imageWriter: DataWriter, + is_debug=False, + start_page_id=0, + end_page_id=None, + lang=None, + *args, + **kwargs +): """ocr和文本混合的pdf，全部解析出来.""" def parse_pdf(method): try: return method( - pdf_bytes, - pdf_models, + dataset, + model_list, imageWriter, start_page_id=start_page_id, end_page_id=end_page_id, @@ -91,12 +115,12 @@ def parse_union_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: DataWriter, pdf_info_dict = parse_pdf(parse_pdf_by_txt) if pdf_info_dict is None or pdf_info_dict.get('_need_drop', False): logger.warning('parse_pdf_by_txt drop or error, switch to parse_pdf_by_ocr') - if input_model_is_empty: + if len(model_list) == 0: layout_model = kwargs.get('layout_model', None) formula_enable = kwargs.get('formula_enable', None) table_enable = kwargs.get('table_enable', None) - pdf_models = doc_analyze( - pdf_bytes, + infer_res = doc_analyze( + dataset, ocr=True, start_page_id=start_page_id, end_page_id=end_page_id, @@ -105,6 +129,7 @@ def parse_union_pdf(pdf_bytes: bytes, pdf_models: list, imageWriter: DataWriter, formula_enable=formula_enable, table_enable=table_enable, ) + model_list = infer_res.get_infer_res() pdf_info_dict = parse_pdf(parse_pdf_by_ocr) if pdf_info_dict is None: raise Exception('Both parse_pdf_by_txt and parse_pdf_by_ocr failed.') From 4a82d6a07ad1137534d4231d114a880586879eb2 Mon Sep 17 00:00:00 2001 From: icecraft Date: Thu, 28 Nov 2024 15:34:04 +0800 Subject: [PATCH 03/50] feat: add function definitions --- magic_pdf/data/dataset.py | 36 ++++ .../model/doc_analyze_by_custom_model.py | 2 +- magic_pdf/model/operators.py | 177 ++++++++++++++++++ magic_pdf/model/types.py | 122 ------------ magic_pdf/pdf_parse_union_core_v2.py | 25 +-- magic_pdf/pipe/operators.py | 124 ++++++++++++ magic_pdf/pipe/types.py | 62 ------ magic_pdf/tools/common.py | 5 +- 8 files changed, 353 insertions(+), 200 deletions(-) create mode 100644 magic_pdf/model/operators.py delete mode 100644 magic_pdf/model/types.py create mode 100644 magic_pdf/pipe/operators.py delete mode 100644 magic_pdf/pipe/types.py diff --git a/magic_pdf/data/dataset.py b/magic_pdf/data/dataset.py index d71d6b84..388d4ce0 100644 --- a/magic_pdf/data/dataset.py +++ b/magic_pdf/data/dataset.py @@ -32,10 +32,28 @@ class PageableData(ABC): @abstractmethod def draw_rect(self, rect_coords, color, fill, fill_opacity, width, overlay): + """draw rectangle. + + Args: + rect_coords (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] + color (list[float] | None): three element tuple which descript the RGB of the board line, None means no board line + fill (list[float] | None): fill the board with RGB, None means will not fill with color + fill_opacity (float): opacity of the fill, range from [0, 1] + width (float): the width of board + overlay (bool): fill the color in foreground or background. True means fill in background. + """ pass @abstractmethod def insert_text(self, coord, content, fontsize, color): + """insert text. + + Args: + coord (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] + content (str): the text content + fontsize (int): font size of the text + color (list[float] | None): three element tuple which descript the RGB of the board line, None will use the default font color! + """ pass @@ -244,6 +262,16 @@ class Doc(PageableData): return getattr(self._doc, name) def draw_rect(self, rect_coords, color, fill, fill_opacity, width, overlay): + """draw rectangle. + + Args: + rect_coords (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] + color (list[float] | None): three element tuple which descript the RGB of the board line, None means no board line + fill (list[float] | None): fill the board with RGB, None means will not fill with color + fill_opacity (float): opacity of the fill, range from [0, 1] + width (float): the width of board + overlay (bool): fill the color in foreground or background. True means fill in background. + """ self._doc.draw_rect( rect_coords, color=color, @@ -254,4 +282,12 @@ class Doc(PageableData): ) def insert_text(self, coord, content, fontsize, color): + """insert text. + + Args: + coord (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] + content (str): the text content + fontsize (int): font size of the text + color (list[float] | None): three element tuple which descript the RGB of the board line, None will use the default font color! + """ self._doc.insert_text(coord, content, fontsize=fontsize, color=color) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index aabc7d2c..17a314ff 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -13,7 +13,7 @@ from magic_pdf.libs.config_reader import (get_device, get_formula_config, get_local_models_dir, get_table_recog_config) from magic_pdf.model.model_list import MODEL -from magic_pdf.model.types import InferenceResult +from magic_pdf.model.operators import InferenceResult def dict_compare(d1, d2): diff --git a/magic_pdf/model/operators.py b/magic_pdf/model/operators.py new file mode 100644 index 00000000..41933400 --- /dev/null +++ b/magic_pdf/model/operators.py @@ -0,0 +1,177 @@ +import copy +import json +import os +from typing import Callable + +from magic_pdf.config.enums import SupportedPdfParseMethod +from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset +from magic_pdf.filter import classify +from magic_pdf.libs.draw_bbox import draw_model_bbox +from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union +from magic_pdf.pipe.operators import PipeResult + + +class InferenceResult: + def __init__(self, inference_results: list, dataset: Dataset): + """Initialized method. + + Args: + inference_results (list): the inference result generated by model + dataset (Dataset): the dataset related with model inference result + """ + self._infer_res = inference_results + self._dataset = dataset + + def draw_model(self, file_path: str) -> None: + """Draw model inference result. + + Args: + file_path (str): the output file path + """ + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + draw_model_bbox( + copy.deepcopy(self._infer_res), self._dataset, dir_name, base_name + ) + + def dump_model(self, writer: DataWriter, file_path: str): + """Dump model inference result to file. + + Args: + writer (DataWriter): writer handle + file_path (str): the location of target file + """ + writer.write_string( + file_path, json.dumps(self._infer_res, ensure_ascii=False, indent=4) + ) + + def get_infer_res(self): + """Get the inference result. + + Returns: + list[dict]: the inference result generated by model + """ + return self._infer_res + + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(inference_result, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + return proc(copy.deepcopy(self._infer_res), *args, **kwargs) + + def pipe_auto_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + """Post-proc the model inference result. + step1: classify the dataset type + step2: based the result of step1, using `pipe_txt_mode` or `pipe_ocr_mode` + + Args: + imageWriter (DataWriter): the image writer handle + start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process + end_page_id (_type_, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + debug_mode (bool, optional): Defaults to False. will dump more log if enabled + lang (_type_, optional): Defaults to None. + + Returns: + PipeResult: the result + """ + + pdf_proc_method = classify(self._dataset.data_bits()) + + if pdf_proc_method == SupportedPdfParseMethod.TXT: + return self.pipe_txt_mode( + imageWriter, start_page_id, end_page_id, debug_mode, lang + ) + else: + return self.pipe_ocr_mode( + imageWriter, start_page_id, end_page_id, debug_mode, lang + ) + + def pipe_txt_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + """Post-proc the model inference result, Extract the text using the + third library, such as `pymupdf` + + Args: + imageWriter (DataWriter): the image writer handle + start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process + end_page_id (_type_, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + debug_mode (bool, optional): Defaults to False. will dump more log if enabled + lang (_type_, optional): Defaults to None. + + Returns: + PipeResult: the result + """ + + def proc(*args, **kwargs) -> PipeResult: + res = pdf_parse_union(*args, **kwargs) + return PipeResult(res, self._dataset) + + return self.apply( + proc, + self._dataset, + imageWriter, + SupportedPdfParseMethod.TXT, + start_page_id=start_page_id, + end_page_id=end_page_id, + debug_mode=debug_mode, + lang=lang, + ) + + def pipe_ocr_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + """Post-proc the model inference result, Extract the text using `OCR` + technical. + + Args: + imageWriter (DataWriter): the image writer handle + start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process + end_page_id (_type_, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + debug_mode (bool, optional): Defaults to False. will dump more log if enabled + lang (_type_, optional): Defaults to None. + + Returns: + PipeResult: the result + """ + + def proc(*args, **kwargs) -> PipeResult: + res = pdf_parse_union(*args, **kwargs) + return PipeResult(res, self._dataset) + + return self.apply( + proc, + self._dataset, + imageWriter, + SupportedPdfParseMethod.TXT, + start_page_id=start_page_id, + end_page_id=end_page_id, + debug_mode=debug_mode, + lang=lang, + ) diff --git a/magic_pdf/model/types.py b/magic_pdf/model/types.py deleted file mode 100644 index 4514d251..00000000 --- a/magic_pdf/model/types.py +++ /dev/null @@ -1,122 +0,0 @@ -import copy -import json -import os -from typing import Callable - -from magic_pdf.config.enums import SupportedPdfParseMethod -from magic_pdf.data.data_reader_writer import DataWriter -from magic_pdf.data.dataset import Dataset -from magic_pdf.filter import classify -from magic_pdf.libs.draw_bbox import draw_model_bbox -from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union -from magic_pdf.pipe.types import PipeResult - - -class InferenceResult: - def __init__(self, inference_results: list, dataset: Dataset): - self._infer_res = inference_results - self._dataset = dataset - - def draw_model(self, file_path: str) -> None: - dir_name = os.path.dirname(file_path) - base_name = os.path.basename(file_path) - if not os.path.exists(dir_name): - os.makedirs(dir_name, exist_ok=True) - draw_model_bbox( - copy.deepcopy(self._infer_res), self._dataset, dir_name, base_name - ) - - def dump_model(self, writer: DataWriter, file_path: str): - writer.write_string( - file_path, json.dumps(self._infer_res, ensure_ascii=False, indent=4) - ) - - def get_infer_res(self): - return self._infer_res - - def apply(self, proc: Callable, *args, **kwargs): - return proc(copy.deepcopy(self._infer_res), *args, **kwargs) - - def pipe_auto_mode( - self, - imageWriter: DataWriter, - start_page_id=0, - end_page_id=None, - debug_mode=False, - lang=None, - ) -> PipeResult: - def proc(*args, **kwargs) -> PipeResult: - res = pdf_parse_union(*args, **kwargs) - return PipeResult(res, self._dataset) - - pdf_proc_method = classify(self._dataset.data_bits()) - - if pdf_proc_method == SupportedPdfParseMethod.TXT: - return self.apply( - proc, - self._dataset, - imageWriter, - SupportedPdfParseMethod.TXT, - start_page_id=0, - end_page_id=None, - debug_mode=False, - lang=None, - ) - else: - return self.apply( - proc, - self._dataset, - imageWriter, - SupportedPdfParseMethod.OCR, - start_page_id=0, - end_page_id=None, - debug_mode=False, - lang=None, - ) - - def pipe_txt_mode( - self, - imageWriter: DataWriter, - start_page_id=0, - end_page_id=None, - debug_mode=False, - lang=None, - ) -> PipeResult: - def proc(*args, **kwargs) -> PipeResult: - res = pdf_parse_union(*args, **kwargs) - return PipeResult(res, self._dataset) - - return self.apply( - proc, - self._dataset, - imageWriter, - SupportedPdfParseMethod.TXT, - start_page_id=0, - end_page_id=None, - debug_mode=False, - lang=None, - ) - - def pipe_ocr_mode( - self, - imageWriter: DataWriter, - start_page_id=0, - end_page_id=None, - debug_mode=False, - lang=None, - ) -> PipeResult: - - def proc(*args, **kwargs) -> PipeResult: - res = pdf_parse_union(*args, **kwargs) - return PipeResult(res, self._dataset) - - return self.apply( - proc, - self._dataset, - imageWriter, - SupportedPdfParseMethod.TXT, - start_page_id=0, - end_page_id=None, - debug_mode=False, - lang=None, - ) diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index bd4b6c6a..a5d6efe4 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -4,8 +4,8 @@ import statistics import time from typing import List -import torch import fitz +import torch from loguru import logger from magic_pdf.config.enums import SupportedPdfParseMethod @@ -16,17 +16,13 @@ from magic_pdf.libs.clean_memory import clean_memory from magic_pdf.libs.config_reader import get_local_layoutreader_model_dir from magic_pdf.libs.convert_utils import dict_to_list from magic_pdf.libs.hash_utils import compute_md5 - from magic_pdf.libs.pdf_image_tools import cut_image_to_pil_image from magic_pdf.model.magic_model import MagicModel -os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新 -os.environ['YOLO_VERBOSE'] = 'False' # disable yolo logger - try: import torchtext - if torchtext.__version__ >= "0.18.0": + if torchtext.__version__ >= '0.18.0': torchtext.disable_torchtext_deprecation_warning() except ImportError: pass @@ -39,6 +35,9 @@ from magic_pdf.pre_proc.ocr_detect_all_bboxes import ocr_prepare_bboxes_for_layo from magic_pdf.pre_proc.ocr_dict_merge import fill_spans_in_blocks, fix_block_spans_v2, fix_discarded_block from magic_pdf.pre_proc.ocr_span_list_modify import get_qa_need_list_v2, remove_overlaps_low_confidence_spans, remove_overlaps_min_spans +os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新 +os.environ['YOLO_VERBOSE'] = 'False' # disable yolo logger + def __replace_STX_ETX(text_str: str): """Replace \u0002 and \u0003, as these characters become garbled when extracted using pymupdf. In fact, they were originally quotation marks. @@ -90,7 +89,10 @@ def chars_to_content(span): LINE_STOP_FLAG = ('.', '!', '?', '。', '！', '？', ')', '）', '"', '”', ':', '：', ';', '；', ']', '】', '}', '}', '>', '》', '、', ',', '，', '-', '—', '–',) +<<<<<<< HEAD LINE_START_FLAG = ('(', '（', '"', '“', '【', '{', '《', '<', '「', '『', '【', '[',) +======= +>>>>>>> 731f4bf (feat: add function definitions) def fill_char_in_spans(spans, all_chars): @@ -233,7 +235,7 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang # 初始化ocr模型 atom_model_manager = AtomModelSingleton() ocr_model = atom_model_manager.get_atom_model( - atom_model_name="ocr", + atom_model_name='ocr', ocr_show_log=False, det_db_box_thresh=0.3, lang=lang @@ -241,7 +243,7 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang for span in empty_spans: # 对span的bbox截图再ocr - span_img = cut_image_to_pil_image(span['bbox'], pdf_page, mode="cv2") + span_img = cut_image_to_pil_image(span['bbox'], pdf_page, mode='cv2') ocr_res = ocr_model.ocr(span_img, det=False) if ocr_res and len(ocr_res) > 0: if len(ocr_res[0]) > 0: @@ -681,7 +683,7 @@ def parse_page_core( """根据parse_mode，构造spans，主要是文本类的字符填充""" if parse_mode == SupportedPdfParseMethod.TXT: - """使用新版本的混合ocr方案""" + """使用新版本的混合ocr方案.""" spans = txt_spans_extract_v2(page_doc, spans, all_bboxes, all_discarded_blocks, lang) elif parse_mode == SupportedPdfParseMethod.OCR: @@ -689,7 +691,6 @@ def parse_page_core( else: raise Exception('parse_mode must be txt or ocr') - """先处理不需要排版的discarded_blocks""" discarded_block_with_spans, spans = fill_spans_in_blocks( all_discarded_blocks, spans, 0.4 @@ -762,8 +763,8 @@ def parse_page_core( def pdf_parse_union( - dataset: Dataset, model_list, + dataset: Dataset, imageWriter, parse_mode, start_page_id=0, @@ -832,4 +833,4 @@ def pdf_parse_union( if __name__ == '__main__': - pass \ No newline at end of file + pass diff --git a/magic_pdf/pipe/operators.py b/magic_pdf/pipe/operators.py new file mode 100644 index 00000000..e5022711 --- /dev/null +++ b/magic_pdf/pipe/operators.py @@ -0,0 +1,124 @@ +import json +import os + +from magic_pdf.config.make_content_config import DropMode, MakeMode +from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset +from magic_pdf.dict2md.ocr_mkcontent import union_make +from magic_pdf.libs.draw_bbox import (draw_layout_bbox, draw_line_sort_bbox, + draw_span_bbox) +from magic_pdf.libs.json_compressor import JsonCompressor + + +class PipeResult: + def __init__(self, pipe_res, dataset: Dataset): + """Initialized. + + Args: + pipe_res (list[dict]): the pipeline processed result of model inference result + dataset (Dataset): the dataset associated with pipe_res + """ + self._pipe_res = pipe_res + self._dataset = dataset + + def dump_md( + self, + writer: DataWriter, + file_path: str, + img_dir_or_bucket_prefix: str, + drop_mode=DropMode.WHOLE_PDF, + md_make_mode=MakeMode.MM_MD, + ): + """Dump The Markdown. + + Args: + writer (DataWriter): File writer handle + file_path (str): The file location of markdown + img_dir_or_bucket_prefix (str): The s3 bucket prefix or local file directory which used to store the figure + drop_mode (str, optional): Drop strategy when some page which is corrupted or inappropriate. Defaults to DropMode.WHOLE_PDF. + md_make_mode (str, optional): The content Type of Markdown be made. Defaults to MakeMode.MM_MD. + """ + pdf_info_list = self._pipe_res['pdf_info'] + md_content = union_make( + pdf_info_list, md_make_mode, drop_mode, img_dir_or_bucket_prefix + ) + writer.write_string(file_path, md_content) + + def dump_content_list( + self, writer: DataWriter, file_path: str, image_dir_or_bucket_prefix: str + ): + """Dump Content List. + + Args: + writer (DataWriter): File writer handle + file_path (str): The file location of content list + image_dir_or_bucket_prefix (str): The s3 bucket prefix or local file directory which used to store the figure + """ + pdf_info_list = self._pipe_res['pdf_info'] + content_list = union_make( + pdf_info_list, + MakeMode.STANDARD_FORMAT, + DropMode.NONE, + image_dir_or_bucket_prefix, + ) + writer.write_string( + file_path, json.dumps(content_list, ensure_ascii=False, indent=4) + ) + + def dump_middle_json(self, writer: DataWriter, file_path: str): + """Dump the result of pipeline. + + Args: + writer (DataWriter): File writer handler + file_path (str): The file location of middle json + """ + writer.write_string( + file_path, json.dumps(self._pipe_res, ensure_ascii=False, indent=4) + ) + + def draw_layout(self, file_path: str) -> None: + """Draw the layout. + + Args: + file_path (str): The file location of layout result file + """ + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + pdf_info = self._pipe_res['pdf_info'] + draw_layout_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) + + def draw_span(self, file_path: str): + """Draw the Span. + + Args: + file_path (str): The file location of span result file + """ + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + pdf_info = self._pipe_res['pdf_info'] + draw_span_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) + + def draw_line_sort(self, file_path: str): + """Draw line sort. + + Args: + file_path (str): The file location of line sort result file + """ + dir_name = os.path.dirname(file_path) + base_name = os.path.basename(file_path) + if not os.path.exists(dir_name): + os.makedirs(dir_name, exist_ok=True) + pdf_info = self._pipe_res['pdf_info'] + draw_line_sort_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) + + def get_compress_pdf_mid_data(self): + """Compress the pipeline result. + + Returns: + str: compress the pipeline result and return + """ + return JsonCompressor.compress_json(self.pdf_mid_data) diff --git a/magic_pdf/pipe/types.py b/magic_pdf/pipe/types.py deleted file mode 100644 index 6ee9b2ce..00000000 --- a/magic_pdf/pipe/types.py +++ /dev/null @@ -1,62 +0,0 @@ - -import json -import os - -from magic_pdf.config.make_content_config import DropMode, MakeMode -from magic_pdf.data.data_reader_writer import DataWriter -from magic_pdf.data.dataset import Dataset -from magic_pdf.dict2md.ocr_mkcontent import union_make -from magic_pdf.libs.draw_bbox import (draw_layout_bbox, draw_line_sort_bbox, - draw_span_bbox) -from magic_pdf.libs.json_compressor import JsonCompressor - - -class PipeResult: - def __init__(self, pipe_res, dataset: Dataset): - self._pipe_res = pipe_res - self._dataset = dataset - - def dump_md(self, writer: DataWriter, file_path: str, img_dir_or_bucket_prefix: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD): - pdf_info_list = self._pipe_res['pdf_info'] - md_content = union_make(pdf_info_list, md_make_mode, drop_mode, img_dir_or_bucket_prefix) - writer.write_string(file_path, md_content) - - def dump_content_list(self, writer: DataWriter, file_path: str, image_dir_or_bucket_prefix: str, drop_mode=DropMode.NONE): - pdf_info_list = self._pipe_res['pdf_info'] - content_list = union_make(pdf_info_list, MakeMode.STANDARD_FORMAT, drop_mode, image_dir_or_bucket_prefix) - writer.write_string(file_path, json.dumps(content_list, ensure_ascii=False, indent=4)) - - def dump_middle_json(self, writer: DataWriter, file_path: str): - writer.write_string(file_path, json.dumps(self._pipe_res, ensure_ascii=False, indent=4)) - - def draw_layout(self, file_path: str) -> None: - dir_name = os.path.dirname(file_path) - base_name = os.path.basename(file_path) - if not os.path.exists(dir_name): - os.makedirs(dir_name, exist_ok=True) - pdf_info = self._pipe_res['pdf_info'] - draw_layout_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) - - def draw_span(self, file_path: str): - dir_name = os.path.dirname(file_path) - base_name = os.path.basename(file_path) - if not os.path.exists(dir_name): - os.makedirs(dir_name, exist_ok=True) - pdf_info = self._pipe_res['pdf_info'] - draw_span_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) - - def draw_line_sort(self, file_path: str): - dir_name = os.path.dirname(file_path) - base_name = os.path.basename(file_path) - if not os.path.exists(dir_name): - os.makedirs(dir_name, exist_ok=True) - pdf_info = self._pipe_res['pdf_info'] - draw_line_sort_bbox(pdf_info, self._dataset.data_bits(), dir_name, base_name) - - def draw_content_list(self, writer: DataWriter, file_path: str, img_dir_or_bucket_prefix: str, drop_mode=DropMode.WHOLE_PDF): - pdf_info_list = self._pipe_res['pdf_info'] - content_list = union_make(pdf_info_list, MakeMode.STANDARD_FORMAT, drop_mode, img_dir_or_bucket_prefix) - writer.write_string(file_path, json.dumps(content_list, ensure_ascii=False, indent=4)) - - def get_compress_pdf_mid_data(self): - return JsonCompressor.compress_json(self.pdf_mid_data) diff --git a/magic_pdf/tools/common.py b/magic_pdf/tools/common.py index 3bd2443b..4f416f40 100644 --- a/magic_pdf/tools/common.py +++ b/magic_pdf/tools/common.py @@ -10,7 +10,7 @@ from magic_pdf.config.make_content_config import DropMode, MakeMode from magic_pdf.data.data_reader_writer import FileBasedDataWriter from magic_pdf.data.dataset import PymuDocDataset from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze -from magic_pdf.model.types import InferenceResult +from magic_pdf.model.operators import InferenceResult # from io import BytesIO # from pypdf import PdfReader, PdfWriter @@ -223,8 +223,7 @@ def do_parse( pipe_result.dump_content_list( md_writer, f'{pdf_file_name}_content_list.json', - image_dir, - drop_mode=DropMode.NONE, + image_dir ) logger.info(f'local output dir is {local_md_dir}') From d44e7a28f422c8f7b790f27d2fd71aabf023f456 Mon Sep 17 00:00:00 2001 From: xu rui Date: Fri, 29 Nov 2024 15:51:59 +0800 Subject: [PATCH 04/50] refactor: add docs --- magic_pdf/model/__init__.py | 124 +++++++++++++++++++++++++++ magic_pdf/model/operators.py | 17 ++-- next_docs/en/api.rst | 2 + next_docs/en/api/model_operators.rst | 8 ++ next_docs/en/api/pipe_operators.rst | 9 ++ next_docs/en/conf.py | 2 +- 6 files changed, 153 insertions(+), 9 deletions(-) create mode 100644 next_docs/en/api/model_operators.rst create mode 100644 next_docs/en/api/pipe_operators.rst diff --git a/magic_pdf/model/__init__.py b/magic_pdf/model/__init__.py index 08653f24..1b0b5b22 100644 --- a/magic_pdf/model/__init__.py +++ b/magic_pdf/model/__init__.py @@ -1,2 +1,126 @@ +from typing import Callable + +from abc import ABC, abstractmethod + +from magic_pdf.data.data_reader_writer import DataWriter +from magic_pdf.data.dataset import Dataset +from magic_pdf.pipe.operators import PipeResult + + __use_inside_model__ = True __model_mode__ = "full" + + +class InferenceResultBase(ABC): + + @abstractmethod + def __init__(self, inference_results: list, dataset: Dataset): + """Initialized method. + + Args: + inference_results (list): the inference result generated by model + dataset (Dataset): the dataset related with model inference result + """ + self._infer_res = inference_results + self._dataset = dataset + + @abstractmethod + def draw_model(self, file_path: str) -> None: + """Draw model inference result. + + Args: + file_path (str): the output file path + """ + pass + + @abstractmethod + def dump_model(self, writer: DataWriter, file_path: str): + """Dump model inference result to file. + + Args: + writer (DataWriter): writer handle + file_path (str): the location of target file + """ + pass + + @abstractmethod + def get_infer_res(self): + """Get the inference result. + + Returns: + list: the inference result generated by model + """ + pass + + @abstractmethod + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(inference_result, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + pass + + @abstractmethod + def pipe_auto_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + """Post-proc the model inference result. + step1: classify the dataset type + step2: based the result of step1, using `pipe_txt_mode` or `pipe_ocr_mode` + + Args: + imageWriter (DataWriter): the image writer handle + start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process + end_page_id (int, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + debug_mode (bool, optional): Defaults to False. will dump more log if enabled + lang (str, optional): Defaults to None. + + Returns: + PipeResult: the result + """ + pass + + @abstractmethod + def pipe_txt_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + """Post-proc the model inference result, Extract the text using the + third library, such as `pymupdf` + + Args: + imageWriter (DataWriter): the image writer handle + start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process + end_page_id (int, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + debug_mode (bool, optional): Defaults to False. will dump more log if enabled + lang (str, optional): Defaults to None. + + Returns: + PipeResult: the result + """ + pass + + @abstractmethod + def pipe_ocr_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + pass diff --git a/magic_pdf/model/operators.py b/magic_pdf/model/operators.py index 41933400..a69cf981 100644 --- a/magic_pdf/model/operators.py +++ b/magic_pdf/model/operators.py @@ -10,9 +10,10 @@ from magic_pdf.filter import classify from magic_pdf.libs.draw_bbox import draw_model_bbox from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union from magic_pdf.pipe.operators import PipeResult +from magic_pdf.model import InferenceResultBase -class InferenceResult: +class InferenceResult(InferenceResultBase): def __init__(self, inference_results: list, dataset: Dataset): """Initialized method. @@ -52,7 +53,7 @@ class InferenceResult: """Get the inference result. Returns: - list[dict]: the inference result generated by model + list: the inference result generated by model """ return self._infer_res @@ -83,9 +84,9 @@ class InferenceResult: Args: imageWriter (DataWriter): the image writer handle start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process - end_page_id (_type_, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + end_page_id (int, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process debug_mode (bool, optional): Defaults to False. will dump more log if enabled - lang (_type_, optional): Defaults to None. + lang (str, optional): Defaults to None. Returns: PipeResult: the result @@ -116,9 +117,9 @@ class InferenceResult: Args: imageWriter (DataWriter): the image writer handle start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process - end_page_id (_type_, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + end_page_id (int, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process debug_mode (bool, optional): Defaults to False. will dump more log if enabled - lang (_type_, optional): Defaults to None. + lang (str, optional): Defaults to None. Returns: PipeResult: the result @@ -153,9 +154,9 @@ class InferenceResult: Args: imageWriter (DataWriter): the image writer handle start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process - end_page_id (_type_, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process + end_page_id (int, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process debug_mode (bool, optional): Defaults to False. will dump more log if enabled - lang (_type_, optional): Defaults to None. + lang (str, optional): Defaults to None. Returns: PipeResult: the result diff --git a/next_docs/en/api.rst b/next_docs/en/api.rst index 787b5fe9..efdbf567 100644 --- a/next_docs/en/api.rst +++ b/next_docs/en/api.rst @@ -7,3 +7,5 @@ api/read_api api/schemas api/io + api/pipe_operators + api/model_operators \ No newline at end of file diff --git a/next_docs/en/api/model_operators.rst b/next_docs/en/api/model_operators.rst new file mode 100644 index 00000000..ccf85128 --- /dev/null +++ b/next_docs/en/api/model_operators.rst @@ -0,0 +1,8 @@ + +Model Api +========== + +.. autoclass:: magic_pdf.model.InferenceResultBase + :members: + :inherited-members: + :show-inheritance: diff --git a/next_docs/en/api/pipe_operators.rst b/next_docs/en/api/pipe_operators.rst new file mode 100644 index 00000000..df25de17 --- /dev/null +++ b/next_docs/en/api/pipe_operators.rst @@ -0,0 +1,9 @@ + + +Pipeline Api +============= + +.. autoclass:: magic_pdf.pipe.operators.PipeResult + :members: + :inherited-members: + :show-inheritance: \ No newline at end of file diff --git a/next_docs/en/conf.py b/next_docs/en/conf.py index 5968eec2..857ef27e 100644 --- a/next_docs/en/conf.py +++ b/next_docs/en/conf.py @@ -114,7 +114,7 @@ autodoc_mock_imports = [ 'sentencepiece', 'vllm.cuda_utils', 'vllm._C', - 'numpy', + # 'numpy', 'tqdm', ] From f6bd47de6a27d5f658d1940312d2807ec4336690 Mon Sep 17 00:00:00 2001 From: xu rui Date: Mon, 2 Dec 2024 15:01:47 +0800 Subject: [PATCH 05/50] docs: add dataset method description --- magic_pdf/data/dataset.py | 86 ++++++++++++- magic_pdf/libs/draw_bbox.py | 1 - magic_pdf/pdf_parse_union_core_v2.py | 3 - magic_pdf/pipe/operators.py | 14 ++ .../en/user_guide/quick_start/to_markdown.rst | 102 ++++++++------- .../user_guide/quick_start/to_markdown.rst | 121 ++++++++++-------- 6 files changed, 217 insertions(+), 110 deletions(-) diff --git a/magic_pdf/data/dataset.py b/magic_pdf/data/dataset.py index 388d4ce0..ce44a712 100644 --- a/magic_pdf/data/dataset.py +++ b/magic_pdf/data/dataset.py @@ -36,7 +36,7 @@ class PageableData(ABC): Args: rect_coords (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] - color (list[float] | None): three element tuple which descript the RGB of the board line, None means no board line + color (list[float] | None): three element tuple which describe the RGB of the board line, None means no board line fill (list[float] | None): fill the board with RGB, None means will not fill with color fill_opacity (float): opacity of the fill, range from [0, 1] width (float): the width of board @@ -52,7 +52,7 @@ class PageableData(ABC): coord (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] content (str): the text content fontsize (int): font size of the text - color (list[float] | None): three element tuple which descript the RGB of the board line, None will use the default font color! + color (list[float] | None): three element tuple which describe the RGB of the board line, None will use the default font color! """ pass @@ -96,14 +96,39 @@ class Dataset(ABC): @abstractmethod def dump_to_file(self, file_path: str): + """Dump the file + + Args: + file_path (str): the file path + """ pass @abstractmethod def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(dataset, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ pass @abstractmethod def classify(self) -> SupportedPdfParseMethod: + """classify the dataset + + Returns: + SupportedPdfParseMethod: _description_ + """ + pass + + @abstractmethod + def clone(self): + """clone this dataset + """ pass @@ -151,18 +176,42 @@ class PymuDocDataset(Dataset): return self._records[page_id] def dump_to_file(self, file_path: str): + """Dump the file + + Args: + file_path (str): the file path + """ + dir_name = os.path.dirname(file_path) if dir_name not in ('', '.', '..'): os.makedirs(dir_name, exist_ok=True) self._raw_fitz.save(file_path) def apply(self, proc: Callable, *args, **kwargs): - new_args = tuple([self] + list(args)) - return proc(*new_args, **kwargs) + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(dataset, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + return proc(self, *args, **kwargs) def classify(self) -> SupportedPdfParseMethod: + """classify the dataset + + Returns: + SupportedPdfParseMethod: _description_ + """ return classify(self._data_bits) + def clone(self): + """clone this dataset + """ + return PymuDocDataset(self._raw_data) + class ImageDataset(Dataset): def __init__(self, bits: bytes): @@ -209,17 +258,40 @@ class ImageDataset(Dataset): return self._records[page_id] def dump_to_file(self, file_path: str): + """Dump the file + + Args: + file_path (str): the file path + """ dir_name = os.path.dirname(file_path) if dir_name not in ('', '.', '..'): os.makedirs(dir_name, exist_ok=True) self._raw_fitz.save(file_path) def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(dataset, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ return proc(self, *args, **kwargs) def classify(self) -> SupportedPdfParseMethod: + """classify the dataset + + Returns: + SupportedPdfParseMethod: _description_ + """ return SupportedPdfParseMethod.OCR + def clone(self): + """clone this dataset + """ + return ImageDataset(self._raw_data) class Doc(PageableData): """Initialized with pymudoc object.""" @@ -228,7 +300,7 @@ class Doc(PageableData): self._doc = doc def get_image(self): - """Return the imge info. + """Return the image info. Returns: dict: { @@ -266,7 +338,7 @@ class Doc(PageableData): Args: rect_coords (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] - color (list[float] | None): three element tuple which descript the RGB of the board line, None means no board line + color (list[float] | None): three element tuple which describe the RGB of the board line, None means no board line fill (list[float] | None): fill the board with RGB, None means will not fill with color fill_opacity (float): opacity of the fill, range from [0, 1] width (float): the width of board @@ -288,6 +360,6 @@ class Doc(PageableData): coord (list[float]): four elements array contain the top-left and bottom-right coordinates, [x0, y0, x1, y1] content (str): the text content fontsize (int): font size of the text - color (list[float] | None): three element tuple which descript the RGB of the board line, None will use the default font color! + color (list[float] | None): three element tuple which describe the RGB of the board line, None will use the default font color! """ self._doc.insert_text(coord, content, fontsize=fontsize, color=color) diff --git a/magic_pdf/libs/draw_bbox.py b/magic_pdf/libs/draw_bbox.py index e325ee45..7ec359ea 100644 --- a/magic_pdf/libs/draw_bbox.py +++ b/magic_pdf/libs/draw_bbox.py @@ -3,7 +3,6 @@ from magic_pdf.config.constants import CROSS_PAGE from magic_pdf.config.ocr_content_type import (BlockType, CategoryId, ContentType) from magic_pdf.data.dataset import Dataset -from magic_pdf.libs.commons import fitz # PyMuPDF from magic_pdf.model.magic_model import MagicModel diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index a5d6efe4..61096999 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -89,10 +89,7 @@ def chars_to_content(span): LINE_STOP_FLAG = ('.', '!', '?', '。', '！', '？', ')', '）', '"', '”', ':', '：', ';', '；', ']', '】', '}', '}', '>', '》', '、', ',', '，', '-', '—', '–',) -<<<<<<< HEAD LINE_START_FLAG = ('(', '（', '"', '“', '【', '{', '《', '<', '「', '『', '【', '[',) -======= ->>>>>>> 731f4bf (feat: add function definitions) def fill_char_in_spans(spans, all_chars): diff --git a/magic_pdf/pipe/operators.py b/magic_pdf/pipe/operators.py index e5022711..f8a6168a 100644 --- a/magic_pdf/pipe/operators.py +++ b/magic_pdf/pipe/operators.py @@ -1,5 +1,7 @@ import json import os +from typing import Callable +import copy from magic_pdf.config.make_content_config import DropMode, MakeMode from magic_pdf.data.data_reader_writer import DataWriter @@ -122,3 +124,15 @@ class PipeResult: str: compress the pipeline result and return """ return JsonCompressor.compress_json(self.pdf_mid_data) + + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(pipeline_result, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + return proc(copy.deepcopy(self._pipe_res), *args, **kwargs) diff --git a/next_docs/en/user_guide/quick_start/to_markdown.rst b/next_docs/en/user_guide/quick_start/to_markdown.rst index 94b9d5c7..fc2c11d2 100644 --- a/next_docs/en/user_guide/quick_start/to_markdown.rst +++ b/next_docs/en/user_guide/quick_start/to_markdown.rst @@ -12,17 +12,17 @@ Local File Example import os from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader - from magic_pdf.config.make_content_config import DropMode, MakeMode - from magic_pdf.pipe.OCRPipe import OCRPipe + from magic_pdf.data.dataset import PymuDocDataset + from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze - - ## args - model_list = [] + # args pdf_file_name = "abc.pdf" # replace with the real pdf path + name_without_suff = pdf_file_name.split(".")[0] - - ## prepare env + # prepare env local_image_dir, local_md_dir = "output/images", "output" + image_dir = str(os.path.basename(local_image_dir)) + os.makedirs(local_image_dir, exist_ok=True) image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter( @@ -30,27 +30,31 @@ Local File Example ) image_dir = str(os.path.basename(local_image_dir)) + # read bytes reader1 = FileBasedDataReader("") - pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) - pipe = OCRPipe(pdf_bytes, model_list, image_writer) + ## inference + infer_result = ds.apply(doc_analyze, ocr=True) - pipe.pipe_classify() - pipe.pipe_analyze() - pipe.pipe_parse() + ### draw model result on each page + infer_result.draw_model(os.path.join(local_md_dir, f"{name_without_suff}_model.pdf")) - pdf_info = pipe.pdf_mid_data["pdf_info"] + ## pipeline + pipe_result = infer_result.pipe_ocr_mode(image_writer) + ### draw layout result on each page + pipe_result.draw_layout(os.path.join(local_md_dir, f"{name_without_suff}_layout.pdf")) - md_content = pipe.pipe_mk_markdown( - image_dir, drop_mode=DropMode.NONE, md_make_mode=MakeMode.MM_MD - ) + ### draw spans result on each page + pipe_result.draw_span(os.path.join(local_md_dir, f"{name_without_suff}_spans.pdf")) - if isinstance(md_content, list): - md_writer.write_string(f"{pdf_file_name}.md", "\n".join(md_content)) - else: - md_writer.write_string(f"{pdf_file_name}.md", md_content) + ### dump markdown + pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir) S3 File Example @@ -61,8 +65,8 @@ S3 File Example import os from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter - from magic_pdf.config.make_content_config import DropMode, MakeMode - from magic_pdf.pipe.OCRPipe import OCRPipe + from magic_pdf.data.dataset import PymuDocDataset + from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze bucket_name = "{Your S3 Bucket Name}" # replace with real bucket name ak = "{Your S3 access key}" # replace with real s3 access key @@ -74,29 +78,39 @@ S3 File Example writer = S3DataWriter('unittest/tmp', bucket_name, ak, sk, endpoint_url) image_writer = S3DataWriter('unittest/tmp/images', bucket_name, ak, sk, endpoint_url) - ## args - model_list = [] - pdf_file_name = f"s3://{bucket_name}/{fake pdf path}" # replace with the real s3 path - - pdf_bytes = reader.read(pdf_file_name) # read the pdf content - - - pipe = OCRPipe(pdf_bytes, model_list, image_writer) - - pipe.pipe_classify() - pipe.pipe_analyze() - pipe.pipe_parse() - - pdf_info = pipe.pdf_mid_data["pdf_info"] - - md_content = pipe.pipe_mk_markdown( - "unittest/tmp/images", drop_mode=DropMode.NONE, md_make_mode=MakeMode.MM_MD + # args + pdf_file_name = ( + "s3://llm-pdf-text-1/unittest/tmp/bug5-11.pdf" # replace with the real s3 path ) - if isinstance(md_content, list): - writer.write_string(f"{pdf_file_name}.md", "\n".join(md_content)) - else: - writer.write_string(f"{pdf_file_name}.md", md_content) + # prepare env + local_dir = "output" + name_without_suff = os.path.basename(pdf_file_name).split(".")[0] + + # read bytes + pdf_bytes = reader.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ## inference + infer_result = ds.apply(doc_analyze, ocr=True) + + ### draw model result on each page + infer_result.draw_model(os.path.join(local_dir, f'{name_without_suff}_model.pdf')) # dump to local + + ## pipeline + pipe_result = infer_result.pipe_ocr_mode(image_writer) + + ### 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 -Check :doc:`../data/data_reader_writer` for more [reader | writer] examples +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 diff --git a/next_docs/zh_cn/user_guide/quick_start/to_markdown.rst b/next_docs/zh_cn/user_guide/quick_start/to_markdown.rst index 05549ff2..bbffa449 100644 --- a/next_docs/zh_cn/user_guide/quick_start/to_markdown.rst +++ b/next_docs/zh_cn/user_guide/quick_start/to_markdown.rst @@ -1,28 +1,26 @@ - 转换为 Markdown 文件 ======================== - 本地文件示例 -^^^^^^^^^^^ +^^^^^^^^^^^^^^^^^^ .. code:: python import os from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader - from magic_pdf.config.make_content_config import DropMode, MakeMode - from magic_pdf.pipe.OCRPipe import OCRPipe + from magic_pdf.data.dataset import PymuDocDataset + from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze - - ## args - model_list = [] + # args pdf_file_name = "abc.pdf" # replace with the real pdf path + name_without_suff = pdf_file_name.split(".")[0] - - ## prepare env + # prepare env local_image_dir, local_md_dir = "output/images", "output" + image_dir = str(os.path.basename(local_image_dir)) + os.makedirs(local_image_dir, exist_ok=True) image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter( @@ -30,39 +28,43 @@ ) image_dir = str(os.path.basename(local_image_dir)) + # read bytes reader1 = FileBasedDataReader("") - pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ## inference + infer_result = ds.apply(doc_analyze, ocr=True) + + ### draw model result on each page + infer_result.draw_model(os.path.join(local_md_dir, f"{name_without_suff}_model.pdf")) + + ## pipeline + pipe_result = infer_result.pipe_ocr_mode(image_writer) + + ### draw layout result on each page + pipe_result.draw_layout(os.path.join(local_md_dir, f"{name_without_suff}_layout.pdf")) + + ### draw spans result on each page + pipe_result.draw_span(os.path.join(local_md_dir, f"{name_without_suff}_spans.pdf")) + + ### dump markdown + pipe_result.dump_md(md_writer, f"{name_without_suff}.md", image_dir) - pipe = OCRPipe(pdf_bytes, model_list, image_writer) - - pipe.pipe_classify() - pipe.pipe_analyze() - pipe.pipe_parse() - - pdf_info = pipe.pdf_mid_data["pdf_info"] - - - md_content = pipe.pipe_mk_markdown( - image_dir, drop_mode=DropMode.NONE, md_make_mode=MakeMode.MM_MD - ) - - if isinstance(md_content, list): - md_writer.write_string(f"{pdf_file_name}.md", "\n".join(md_content)) - else: - md_writer.write_string(f"{pdf_file_name}.md", md_content) - - -对象存储使用示例 -^^^^^^^^^^^^^^^ +对象存储文件示例 +^^^^^^^^^^^^^^^^ .. code:: python import os from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter - from magic_pdf.config.make_content_config import DropMode, MakeMode - from magic_pdf.pipe.OCRPipe import OCRPipe + from magic_pdf.data.dataset import PymuDocDataset + from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze bucket_name = "{Your S3 Bucket Name}" # replace with real bucket name ak = "{Your S3 access key}" # replace with real s3 access key @@ -74,30 +76,39 @@ writer = S3DataWriter('unittest/tmp', bucket_name, ak, sk, endpoint_url) image_writer = S3DataWriter('unittest/tmp/images', bucket_name, ak, sk, endpoint_url) - ## args - model_list = [] - pdf_file_name = f"s3://{bucket_name}/{fake pdf path}" # replace with the real s3 path - - pdf_bytes = reader.read(pdf_file_name) # read the pdf content - - - pipe = OCRPipe(pdf_bytes, model_list, image_writer) - - pipe.pipe_classify() - pipe.pipe_analyze() - pipe.pipe_parse() - - pdf_info = pipe.pdf_mid_data["pdf_info"] - - md_content = pipe.pipe_mk_markdown( - "unittest/tmp/images", drop_mode=DropMode.NONE, md_make_mode=MakeMode.MM_MD + # args + pdf_file_name = ( + "s3://llm-pdf-text-1/unittest/tmp/bug5-11.pdf" # replace with the real s3 path ) - if isinstance(md_content, list): - writer.write_string(f"{pdf_file_name}.md", "\n".join(md_content)) - else: - writer.write_string(f"{pdf_file_name}.md", md_content) + # prepare env + local_dir = "output" + name_without_suff = os.path.basename(pdf_file_name).split(".")[0] + # read bytes + pdf_bytes = reader.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ## inference + infer_result = ds.apply(doc_analyze, ocr=True) + + ### draw model result on each page + infer_result.draw_model(os.path.join(local_dir, f'{name_without_suff}_model.pdf')) # dump to local + + ## pipeline + pipe_result = infer_result.pipe_ocr_mode(image_writer) + + ### 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 前去 :doc:`../data/data_reader_writer` 获取更多有关 **读写** 示例 From 91d825b29a8d71f41be094cfa4826f932a210fa4 Mon Sep 17 00:00:00 2001 From: xu rui Date: Mon, 2 Dec 2024 16:02:15 +0800 Subject: [PATCH 06/50] docs: fix table format --- next_docs/requirements.txt | 6 +- .../tutorial/output_file_description.rst | 125 +++++++++--------- 2 files changed, 64 insertions(+), 67 deletions(-) diff --git a/next_docs/requirements.txt b/next_docs/requirements.txt index 1df5f63d..3d4b74bb 100644 --- a/next_docs/requirements.txt +++ b/next_docs/requirements.txt @@ -1,3 +1,7 @@ +numpy==1.26.4 +click==8.1.7 +fast-langdetect==0.2.2 +Brotli==1.1.0 boto3>=1.28.43 loguru>=0.6.0 myst-parser @@ -9,4 +13,4 @@ sphinx-argparse>=0.5.2 sphinx-book-theme>=1.1.3 sphinx-copybutton>=0.5.2 sphinx_rtd_theme>=3.0.1 -autodoc_pydantic>=2.2.0 \ No newline at end of file +autodoc_pydantic>=2.2.0 diff --git a/next_docs/zh_cn/user_guide/tutorial/output_file_description.rst b/next_docs/zh_cn/user_guide/tutorial/output_file_description.rst index 5fcb6d1e..9618e286 100644 --- a/next_docs/zh_cn/user_guide/tutorial/output_file_description.rst +++ b/next_docs/zh_cn/user_guide/tutorial/output_file_description.rst @@ -137,49 +137,45 @@ poly 坐标的格式 [x0, y0, x1, y1, x2, y2, x3, y3], some_pdf_middle.json ~~~~~~~~~~~~~~~~~~~~ -+-----------+----------------------------------------------------------+ -| 字段名 | 解释 | -+===========+==========================================================+ -| pdf_info | list，每个 | -| | 元素都是一个dict,这个dict是每一页pdf的解析结果，详见下表 | -+-----------+----------------------------------------------------------+ -| | ocr \| txt，用来标识本次解析的中间态使用的模式 | -| \_parse_type | | -+-----------+----------------------------------------------------------+ -| | string, 表示本次解析使用的 magic-pdf 的版本号 | -| \_version_name | | -+-----------+----------------------------------------------------------+ ++--------------------+----------------------------------------------------------+ +| 字段名 | 解释 | ++====================+==========================================================+ +| pdf_info | list，每个元素都是一个 | +| | dict，这个dict是每一页pdf的解析结果，详见下表 | ++--------------------+----------------------------------------------------------+ +| \_parse_type | ocr \| txt，用来标识本次解析的中间态使用的模式 | ++--------------------+----------------------------------------------------------+ +| \_version_name | string，表示本次解析使用的 magic-pdf 的版本号 | ++-------------------------------------------------------------------------------+ **pdf_info** 字段结构说明 -+--------------+-------------------------------------------------------+ -| 字段名 | 解释 | -+==============+=======================================================+ -| | pdf预处理后，未分段的中间结果 | -| preeproc_blocks | | -+--------------+-------------------------------------------------------+ -| | 布局分割的结果， | -| layout_bboxes | 含有布局的方向（垂直、水平），和bbox，按阅读顺序排序 | -+--------------+-------------------------------------------------------+ -| page_idx | 页码，从0开始 | -+--------------+-------------------------------------------------------+ -| page_size | 页面的宽度和高度 | -+--------------+-------------------------------------------------------+ -| \ | 布局树状结构 | -| _layout_tree | | -+--------------+-------------------------------------------------------+ -| images | list，每个元素是一个dict，每个dict表示一个img_block | -+--------------+-------------------------------------------------------+ -| tables | list，每个元素是一个dict，每个dict表示一个table_block | -+--------------+-------------------------------------------------------+ -| | list，每个元素 | -| interline_equations | 是一个dict，每个dict表示一个interline_equation_block | -+--------------+-------------------------------------------------------+ -| | List, 模型返回的需要drop的block信息 | -| discarded_blocks | | -+--------------+-------------------------------------------------------+ -| para_blocks | 将preproc_blocks进行分段之后的结果 | -+--------------+-------------------------------------------------------+ ++---------------------+-------------------------------------------------------+ +| 字段名 | 解释 | ++=====================+=======================================================+ +| preproc_blocks | pdf预处理后，未分段的中间结果 | ++---------------------+-------------------------------------------------------+ +| | 布局分割的结果， | +| layout_bboxes | 含有布局的方向（垂直、水平），和bbox，按阅读顺序排序 | ++---------------------+-------------------------------------------------------+ +| page_idx | 页码，从0开始 | ++---------------------+-------------------------------------------------------+ +| page_size | 页面的宽度和高度 | ++---------------------+-------------------------------------------------------+ +| \_layout_tree | 布局树状结构 | ++---------------------+-------------------------------------------------------+ +| images | list，每个元素是一个dict，每个dict表示一个img_block | ++---------------------+-------------------------------------------------------+ +| tables | list，每个元素是一个dict，每个dict表示一个table_block | ++---------------------+-------------------------------------------------------+ +| | list，每个元素是一个 | +| interline_equations | dict，每个dict表示一个interline_equation_block | ++---------------------+-------------------------------------------------------+ +| | List, 模型返回的需要drop的block信息 | +| discarded_blocks | | ++---------------------+-------------------------------------------------------+ +| para_blocks | 将preproc_blocks进行分段之后的结果 | ++---------------------+-------------------------------------------------------+ 上表中 ``para_blocks`` 是个dict的数组，每个dict是一个block结构，block最多支持一次嵌套 @@ -200,20 +196,18 @@ blocks list，里面的每个元素都是一个dict格式的二级block 二级block中的字段包括 -+-----+----------------------------------------------------------------+ -| 字 | 解释 | -| 段 | | -| 名 | | -+=====+================================================================+ -| | block类型 | -| type | | -+-----+----------------------------------------------------------------+ -| | block矩形框坐标 | -| bbox | | -+-----+----------------------------------------------------------------+ -| | list，每个元素都是一个dict表示的line，用来描述一行信息的构成 | -| lines | | -+-----+----------------------------------------------------------------+ ++----------+----------------------------------------------------------------+ +| 字 | 解释 | +| 段 | | +| 名 | | ++==========+================================================================+ +| | block类型 | +| type | | ++----------+----------------------------------------------------------------+ +| bbox | block矩形框坐标 | ++----------+----------------------------------------------------------------+ +| lines | list，每个元素都是一个dict表示的line，用来描述一行信息的构成 | ++----------+----------------------------------------------------------------+ 二级block的类型详解 @@ -237,22 +231,21 @@ interline_equation 行间公式块 line 的 字段格式如下 -+----+-----------------------------------------------------------------+ -| 字 | 解释 | -| 段 | | -| 名 | | -+====+=================================================================+ -| bbox | line的矩形框坐标 | -| | | -+----+-----------------------------------------------------------------+ -| spans | list， | -| | 每个元素都是一个dict表示的span，用来描述一个最小组成单元的构成 | -+----+-----------------------------------------------------------------+ ++-----------+-----------------------------------------------------------------+ +| 字 | 解释 | +| 段 | | +| 名 | | ++===========+=================================================================+ +| bbox | line的矩形框坐标 | ++-----------+-----------------------------------------------------------------+ +| spans | list， | +| | 每个元素都是一个dict表示的span，用来描述一个最小组成单元的构成 | ++-----------+-----------------------------------------------------------------+ **span** +------------+---------------------------------------------------------+ -| 字段名 | 解释 | +| 字段名 | 解释 | +============+=========================================================+ | bbox | span的矩形框坐标 | +------------+---------------------------------------------------------+ From 11994506e0c1d65f6c70be86fe06be2f87b7ece9 Mon Sep 17 00:00:00 2001 From: xu rui Date: Mon, 2 Dec 2024 17:33:03 +0800 Subject: [PATCH 07/50] feat: add zh_cn docs --- magic_pdf/data/dataset.py | 2 +- .../en/_static/image/pipeline.drawio.svg | 4 + .../zh_cn/_static/image/pipeline.drawio.svg | 4 + next_docs/zh_cn/user_guide/tutorial.rst | 2 + .../zh_cn/user_guide/tutorial/pipeline.rst | 175 ++++++++++++++++++ 5 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 next_docs/en/_static/image/pipeline.drawio.svg create mode 100644 next_docs/zh_cn/_static/image/pipeline.drawio.svg create mode 100644 next_docs/zh_cn/user_guide/tutorial/pipeline.rst diff --git a/magic_pdf/data/dataset.py b/magic_pdf/data/dataset.py index ce44a712..bc490d60 100644 --- a/magic_pdf/data/dataset.py +++ b/magic_pdf/data/dataset.py @@ -109,7 +109,7 @@ class Dataset(ABC): Args: proc (Callable): invoke proc as follows: - proc(dataset, *args, **kwargs) + proc(self, *args, **kwargs) Returns: Any: return the result generated by proc diff --git a/next_docs/en/_static/image/pipeline.drawio.svg b/next_docs/en/_static/image/pipeline.drawio.svg new file mode 100644 index 00000000..b7466643 --- /dev/null +++ b/next_docs/en/_static/image/pipeline.drawio.svg @@ -0,0 +1,4 @@ + + + +
Dataset
Inference
Result
Pipeline
Result
/data/abc.pdf
/data/a.jpg

s3://data/a.pdf
draw layout result
dump markdown file

draw span result
draw model result
dump inference result

doc_analyze
pipe_ocr_mode
 \ No newline at end of file diff --git a/next_docs/zh_cn/_static/image/pipeline.drawio.svg b/next_docs/zh_cn/_static/image/pipeline.drawio.svg new file mode 100644 index 00000000..b7466643 --- /dev/null +++ b/next_docs/zh_cn/_static/image/pipeline.drawio.svg @@ -0,0 +1,4 @@ + + + +
Dataset
Inference
Result
Pipeline
Result
/data/abc.pdf
/data/a.jpg

s3://data/a.pdf
draw layout result
dump markdown file

draw span result
draw model result
dump inference result

doc_analyze
pipe_ocr_mode
 \ No newline at end of file diff --git a/next_docs/zh_cn/user_guide/tutorial.rst b/next_docs/zh_cn/user_guide/tutorial.rst index a910ad7f..9beeb6c3 100644 --- a/next_docs/zh_cn/user_guide/tutorial.rst +++ b/next_docs/zh_cn/user_guide/tutorial.rst @@ -9,3 +9,5 @@ :caption: 教程 tutorial/output_file_description + tutorial/pipeline + diff --git a/next_docs/zh_cn/user_guide/tutorial/pipeline.rst b/next_docs/zh_cn/user_guide/tutorial/pipeline.rst new file mode 100644 index 00000000..9f301368 --- /dev/null +++ b/next_docs/zh_cn/user_guide/tutorial/pipeline.rst @@ -0,0 +1,175 @@ + +流水线管道 +=========== + + +极简示例 +^^^^^^^^ + +.. code:: python + + import os + + 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 + + # args + pdf_file_name = "abc.pdf" # replace with the real pdf path + name_without_suff = pdf_file_name.split(".")[0] + + # prepare env + local_image_dir, local_md_dir = "output/images", "output" + image_dir = str(os.path.basename(local_image_dir)) + + os.makedirs(local_image_dir, exist_ok=True) + + image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter( + local_md_dir + ) + image_dir = str(os.path.basename(local_image_dir)) + + # read bytes + reader1 = FileBasedDataReader("") + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) + + +运行以上的代码，会得到如下的结果 + +.. code:: bash + + output/ + ├── abc.md + └── images + + +除去初始化环境，如建立目录、导入依赖库等逻辑。真正执行将 `pdf` 转换为 `markdown` 的代码片段如下 + +.. code:: + + # read bytes + reader1 = FileBasedDataReader("") + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) + + +`ds.apply(doc_analyze, ocr=True)` 会生成 `InferenceResult` 对象。 `InferenceResult` 的对象执行 `pipe_ocr_mode` 方法会生成 `PipeResult` 对象。 +`PipeResult` 对象执行 `dump_md` 会在指定位置生成 `markdown` 文件。 + + +pipeline 的执行过程如下图所示 + +.. image:: ../../_static/image/pipeline.drawio.svg + + + + +.. admonition:: Tip + :class: tip + + 要想获得更多有关 Dataset、InferenceResult、PipeResult 的使用示例子，请前往 :doc:`../quick_start/to_markdown` + + 要想获得更多有关 Dataset、InferenceResult、PipeResult 的细节信息请前往英文版 MinerU 文档进行查看! + + + +管道组合 +^^^^^^^^^ + +.. code:: python + + class Dataset(ABC): + @abstractmethod + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(self, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + pass + + class InferenceResult(InferenceResultBase): + + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(inference_result, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + return proc(copy.deepcopy(self._infer_res), *args, **kwargs) + + def pipe_ocr_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + pass + + class PipeResult: + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(pipeline_result, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + return proc(copy.deepcopy(self._pipe_res), *args, **kwargs) + +`Dataset` 、 `InferenceResult` 和 `PipeResult` 类均有 `apply` method。可用于组合不同阶段的运算过程。 +如下所示，`MinerU` 提供一套组合这些类的计算过程。 + +.. code:: python + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) + +用户可以根据的需求，自行实现一些组合用的函数。比如用户通过 `apply` 方法实现一个统计 `pdf` 文件页数的功能。 + +.. code:: python + + from magic_pdf.data.data_reader_writer import FileBasedDataReader + from magic_pdf.data.dataset import PymuDocDataset + + # args + pdf_file_name = "abc.pdf" # replace with the real pdf path + + # read bytes + reader1 = FileBasedDataReader("") + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + def count_page(ds)-> int: + return len(ds) + + print("page number: ", ds.apply(count_page)) # will output the page count of `abc.pdf` From e4ed6023e4f480f92eef19dd985b6a23fdd8885f Mon Sep 17 00:00:00 2001 From: xu rui Date: Tue, 3 Dec 2024 14:06:42 +0800 Subject: [PATCH 08/50] docs: add [en|zh_cn] docs --- next_docs/en/user_guide/tutorial.rst | 4 +- next_docs/en/user_guide/tutorial/pipeline.rst | 185 ++++++++++++++++++ .../zh_cn/user_guide/tutorial/pipeline.rst | 16 +- 3 files changed, 198 insertions(+), 7 deletions(-) create mode 100644 next_docs/en/user_guide/tutorial/pipeline.rst diff --git a/next_docs/en/user_guide/tutorial.rst b/next_docs/en/user_guide/tutorial.rst index 20764701..eaed04fd 100644 --- a/next_docs/en/user_guide/tutorial.rst +++ b/next_docs/en/user_guide/tutorial.rst @@ -7,4 +7,6 @@ From the beginning to the end, Show how to using mineru via a minimal project .. toctree:: :maxdepth: 1 - tutorial/output_file_description \ No newline at end of file + tutorial/output_file_description + tutorial/pipeline + diff --git a/next_docs/en/user_guide/tutorial/pipeline.rst b/next_docs/en/user_guide/tutorial/pipeline.rst new file mode 100644 index 00000000..8e73a3f5 --- /dev/null +++ b/next_docs/en/user_guide/tutorial/pipeline.rst @@ -0,0 +1,185 @@ + + +Pipeline +========== + + +Minimal Example +^^^^^^^^^^^^^^^^^ + +.. code:: python + + import os + + 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 + + # args + pdf_file_name = "abc.pdf" # replace with the real pdf path + name_without_suff = pdf_file_name.split(".")[0] + + # prepare env + local_image_dir, local_md_dir = "output/images", "output" + image_dir = str(os.path.basename(local_image_dir)) + + os.makedirs(local_image_dir, exist_ok=True) + + image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter( + local_md_dir + ) + image_dir = str(os.path.basename(local_image_dir)) + + # read bytes + reader1 = FileBasedDataReader("") + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) + +Running the above code will result in the following + + +.. code:: bash + + output/ + ├── abc.md + └── images + + +Excluding the setup of the environment, such as creating directories and importing dependencies, the actual code snippet for converting pdf to markdown is as follows + + +.. code:: python + + # read bytes + reader1 = FileBasedDataReader("") + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) + +``ds.apply(doc_analyze, ocr=True)`` generates an ``InferenceResult`` object. The ``InferenceResult`` object, when executing the ``pipe_ocr_mode`` method, produces a ``PipeResult`` object. +The ``PipeResult`` object, upon executing ``dump_md``, generates a ``markdown`` file at the specified location. + + +The pipeline execution process is illustrated in the following diagram + + +.. image:: ../../_static/image/pipeline.drawio.svg + +.. raw:: html + +

+ +Currently, the process is divided into three stages: data, inference, and processing, which correspond to the ``Dataset``, ``InferenceResult``, and ``PipeResult`` entities in the diagram. +These stages are linked together through methods like ``apply``, ``doc_analyze``, or ``pipe_ocr_mode`` + + +.. admonition:: Tip + :class: tip + + For more examples on how to use ``Dataset``, ``InferenceResult``, and ``PipeResult``, please refer to :doc:`../quick_start/to_markdown` + + For more detailed information about ``Dataset``, ``InferenceResult``, and ``PipeResult``, please refer to :doc:`../../api/dataset`, :doc:`../../api/model_operators`, :doc:`../../api/pipe_operators` + + +Pipeline Composition +^^^^^^^^^^^^^^^^^^^^^ + +.. code:: python + + class Dataset(ABC): + @abstractmethod + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(self, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + pass + + class InferenceResult(InferenceResultBase): + + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(inference_result, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + return proc(copy.deepcopy(self._infer_res), *args, **kwargs) + + def pipe_ocr_mode( + self, + imageWriter: DataWriter, + start_page_id=0, + end_page_id=None, + debug_mode=False, + lang=None, + ) -> PipeResult: + pass + + class PipeResult: + def apply(self, proc: Callable, *args, **kwargs): + """Apply callable method which. + + Args: + proc (Callable): invoke proc as follows: + proc(pipeline_result, *args, **kwargs) + + Returns: + Any: return the result generated by proc + """ + return proc(copy.deepcopy(self._pipe_res), *args, **kwargs) + + +The ``Dataset``, ``InferenceResult``, and ``PipeResult`` classes all have an ``apply`` method, which can be used to chain different stages of the computation. +As shown below, ``MinerU`` provides a set of methods to compose these classes. + + +.. code:: python + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) + + +Users can implement their own functions for chaining as needed. For example, a user could use the ``apply`` method to create a function that counts the number of pages in a ``pdf`` file. + + +.. code:: python + + from magic_pdf.data.data_reader_writer import FileBasedDataReader + from magic_pdf.data.dataset import PymuDocDataset + + # args + pdf_file_name = "abc.pdf" # replace with the real pdf path + + # read bytes + reader1 = FileBasedDataReader("") + pdf_bytes = reader1.read(pdf_file_name) # read the pdf content + + # proc + ## Create Dataset Instance + ds = PymuDocDataset(pdf_bytes) + + def count_page(ds)-> int: + return len(ds) + + print("page number: ", ds.apply(count_page)) # will output the page count of `abc.pdf` diff --git a/next_docs/zh_cn/user_guide/tutorial/pipeline.rst b/next_docs/zh_cn/user_guide/tutorial/pipeline.rst index 9f301368..c8adaa3e 100644 --- a/next_docs/zh_cn/user_guide/tutorial/pipeline.rst +++ b/next_docs/zh_cn/user_guide/tutorial/pipeline.rst @@ -49,7 +49,7 @@ └── images -除去初始化环境，如建立目录、导入依赖库等逻辑。真正执行将 `pdf` 转换为 `markdown` 的代码片段如下 +除去初始化环境，如建立目录、导入依赖库等逻辑。真正将 ``pdf`` 转换为 ``markdown`` 的代码片段如下 .. code:: @@ -64,15 +64,19 @@ ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) -`ds.apply(doc_analyze, ocr=True)` 会生成 `InferenceResult` 对象。 `InferenceResult` 的对象执行 `pipe_ocr_mode` 方法会生成 `PipeResult` 对象。 -`PipeResult` 对象执行 `dump_md` 会在指定位置生成 `markdown` 文件。 +``ds.apply(doc_analyze, ocr=True)`` 会生成 ``InferenceResult`` 对象。 ``InferenceResult`` 对象执行 ``pipe_ocr_mode`` 方法会生成 ``PipeResult`` 对象。 +``PipeResult`` 对象执行 ``dump_md`` 会在指定位置生成 ``markdown`` 文件。 pipeline 的执行过程如下图所示 .. image:: ../../_static/image/pipeline.drawio.svg +.. raw:: html +

+ +目前划分出数据、推理、程序处理三个阶段，分别对应着图上的 ``Dataset``， ``InferenceResult``， ``PipeResult`` 这三个实体。通过 ``apply`` ， ``doc_analyze`` 或 ``pipe_ocr_mode`` 等方法链接在一起。 .. admonition:: Tip @@ -140,8 +144,8 @@ pipeline 的执行过程如下图所示 """ return proc(copy.deepcopy(self._pipe_res), *args, **kwargs) -`Dataset` 、 `InferenceResult` 和 `PipeResult` 类均有 `apply` method。可用于组合不同阶段的运算过程。 -如下所示，`MinerU` 提供一套组合这些类的计算过程。 +``Dataset`` 、 ``InferenceResult`` 和 ``PipeResult`` 类均有 ``apply`` method。可用于组合不同阶段的运算过程。 +如下所示，``MinerU`` 提供一套组合这些类的计算过程。 .. code:: python @@ -151,7 +155,7 @@ pipeline 的执行过程如下图所示 ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(md_writer, f"{name_without_suff}.md", image_dir) -用户可以根据的需求，自行实现一些组合用的函数。比如用户通过 `apply` 方法实现一个统计 `pdf` 文件页数的功能。 +用户可以根据的需求，自行实现一些组合用的函数。比如用户通过 ``apply`` 方法实现一个统计 ``pdf`` 文件页数的功能。 .. code:: python From 104273cc79ecc921020e962238014e2ce2dad9ba Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 3 Dec 2024 14:22:35 +0800 Subject: [PATCH 09/50] fix(vram): improve VRAM checking logic - Update VRAM checking logic in app.py and model_utils.py - Add None and type checks for VRAM values - Adjust concurrency limit calculation in app.py - Modify clean_vram function to handle cases with no VRAM information --- magic_pdf/model/sub_modules/model_utils.py | 4 ++-- projects/gradio_app/app.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/magic_pdf/model/sub_modules/model_utils.py b/magic_pdf/model/sub_modules/model_utils.py index c57f50fa..641fd062 100644 --- a/magic_pdf/model/sub_modules/model_utils.py +++ b/magic_pdf/model/sub_modules/model_utils.py @@ -43,7 +43,7 @@ def get_res_list_from_layout_res(layout_res): def clean_vram(device, vram_threshold=8): total_memory = get_vram(device) - if total_memory <= vram_threshold: + if total_memory and total_memory <= vram_threshold: gc_start = time.time() clean_memory() gc_time = round(time.time() - gc_start, 2) @@ -54,4 +54,4 @@ def get_vram(device): if torch.cuda.is_available() and device != 'cpu': total_memory = torch.cuda.get_device_properties(device).total_memory / (1024 ** 3) # 将字节转换为 GB return total_memory - return 0 \ No newline at end of file + return None \ No newline at end of file diff --git a/projects/gradio_app/app.py b/projects/gradio_app/app.py index b4436021..1dffdafb 100644 --- a/projects/gradio_app/app.py +++ b/projects/gradio_app/app.py @@ -187,8 +187,9 @@ def to_pdf(file_path): def get_concurrency_limit(vram_threshold=7.5): vram = get_vram(device = get_device()) - concurrency_limit = int(vram // vram_threshold) - if concurrency_limit < 1: + if vram is not None and isinstance(vram, (int, float)): + concurrency_limit = max(1, int(vram // vram_threshold)) + else: concurrency_limit = 1 # logger.info(f'concurrency_limit: {concurrency_limit}') return concurrency_limit From 044780951ab180961b4f878e20f91f0d128691b7 Mon Sep 17 00:00:00 2001 From: myhloli Date: Thu, 5 Dec 2024 11:09:39 +0800 Subject: [PATCH 10/50] perf(model): add threading lock for OCR model initialization - Introduce a lock to synchronize access to OCR model initialization- This change improves thread safety when multiple threads access the OCR model concurrently - The lock ensures that the OCR model is initialized only once, even in multi-threaded scenarios --- magic_pdf/model/sub_modules/model_init.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/magic_pdf/model/sub_modules/model_init.py b/magic_pdf/model/sub_modules/model_init.py index b9d4d089..0c4b0891 100644 --- a/magic_pdf/model/sub_modules/model_init.py +++ b/magic_pdf/model/sub_modules/model_init.py @@ -82,9 +82,12 @@ def ocr_model_init(show_log: bool = False, return model +from threading import Lock + class AtomModelSingleton: _instance = None _models = {} + _lock = Lock() def __new__(cls, *args, **kwargs): if cls._instance is None: @@ -95,8 +98,13 @@ class AtomModelSingleton: lang = kwargs.get('lang', None) layout_model_name = kwargs.get('layout_model_name', None) key = (atom_model_name, layout_model_name, lang) - if key not in self._models: - self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) + if atom_model_name == AtomicModel.OCR: + with self._lock: + if key not in self._models: + self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) + else: + if key not in self._models: + self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) return self._models[key] From ce592f8bcf2b261d96034993bc25a66e90006e99 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 09:59:22 +0800 Subject: [PATCH 11/50] perf(model): optimize model initialization - Add condition to return existing model if already initialized - Improve efficiency by avoiding redundant model creation --- magic_pdf/model/sub_modules/model_init.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/magic_pdf/model/sub_modules/model_init.py b/magic_pdf/model/sub_modules/model_init.py index 0c4b0891..f4e94bee 100644 --- a/magic_pdf/model/sub_modules/model_init.py +++ b/magic_pdf/model/sub_modules/model_init.py @@ -102,10 +102,13 @@ class AtomModelSingleton: with self._lock: if key not in self._models: self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) + else: + return self._models[key] else: if key not in self._models: self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) - return self._models[key] + else: + return self._models[key] def atom_model_init(model_name: str, **kwargs): From 878f3de00447190aa8348a2f45105305b3881a0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B5=B5=E5=B0=8F=E8=92=99?= Date: Fri, 6 Dec 2024 10:28:25 +0800 Subject: [PATCH 12/50] refactor(magic_pdf): optimize model initialization and threading - Remove unnecessary threading.Lock in AtomModelSingleton - Add threading.Lock to CustomPEKModel for OCR processing - Simplify model initialization logic in AtomModelSingleton --- magic_pdf/model/pdf_extract_kit.py | 12 ++++++++---- magic_pdf/model/sub_modules/model_init.py | 16 +++------------- 2 files changed, 11 insertions(+), 17 deletions(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 2fbcf0a5..78d1930c 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -28,6 +28,8 @@ from magic_pdf.model.sub_modules.model_utils import ( from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import ( get_adjusted_mfdetrec_res, get_ocr_result_list) +from threading import Lock + class CustomPEKModel: @@ -209,16 +211,18 @@ class CustomPEKModel: # ocr识别 ocr_start = time.time() # Process each area that requires OCR processing + lock = Lock() for res in ocr_res_list: new_image, useful_list = crop_img(res, pil_img, crop_paste_x=50, crop_paste_y=50) adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(single_page_mfdetrec_res, useful_list) # OCR recognition new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR) - if self.apply_ocr: - ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] - else: - ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res, rec=False)[0] + with lock: + if self.apply_ocr: + ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] + else: + ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res, rec=False)[0] # Integration results if ocr_res: diff --git a/magic_pdf/model/sub_modules/model_init.py b/magic_pdf/model/sub_modules/model_init.py index f4e94bee..4d4fa1dc 100644 --- a/magic_pdf/model/sub_modules/model_init.py +++ b/magic_pdf/model/sub_modules/model_init.py @@ -82,12 +82,9 @@ def ocr_model_init(show_log: bool = False, return model -from threading import Lock - class AtomModelSingleton: _instance = None _models = {} - _lock = Lock() def __new__(cls, *args, **kwargs): if cls._instance is None: @@ -98,17 +95,10 @@ class AtomModelSingleton: lang = kwargs.get('lang', None) layout_model_name = kwargs.get('layout_model_name', None) key = (atom_model_name, layout_model_name, lang) - if atom_model_name == AtomicModel.OCR: - with self._lock: - if key not in self._models: - self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) - else: - return self._models[key] + if key not in self._models: + self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) else: - if key not in self._models: - self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) - else: - return self._models[key] + return self._models[key] def atom_model_init(model_name: str, **kwargs): From 494859c59eacff24f2ddb03b87c060012f922acf Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 11:27:33 +0800 Subject: [PATCH 13/50] fix: update notify --- .github/workflows/cli.yml | 2 +- .github/workflows/daily.yml | 1 - .github/workflows/huigui.yml | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) delete mode 100644 .github/workflows/daily.yml diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 1ac12e87..0bc08f29 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -37,7 +37,7 @@ jobs: cd $GITHUB_WORKSPACE && pytest -m P0 -s -v tests/test_cli/test_cli_sdk.py notify_to_feishu: - if: ${{ always() && !cancelled() && contains(needs.*.result, 'failure') && (github.ref_name == 'master') }} + if: ${{ always() && !cancelled() && contains(needs.*.result, 'failure')}} needs: cli-test runs-on: pdf steps: diff --git a/.github/workflows/daily.yml b/.github/workflows/daily.yml deleted file mode 100644 index 8b137891..00000000 --- a/.github/workflows/daily.yml +++ /dev/null @@ -1 +0,0 @@ - diff --git a/.github/workflows/huigui.yml b/.github/workflows/huigui.yml index 428f7913..e45fa700 100644 --- a/.github/workflows/huigui.yml +++ b/.github/workflows/huigui.yml @@ -36,7 +36,7 @@ jobs: cd $GITHUB_WORKSPACE && pytest -s -v tests/test_cli/test_cli_sdk.py notify_to_feishu: - if: ${{ always() && !cancelled() && contains(needs.*.result, 'failure') && (github.ref_name == 'master') }} + if: ${{ always() && !cancelled() && contains(needs.*.result, 'failure')}} needs: cli-test runs-on: pdf steps: From eb021e53cca193a652dd56ffc80e2ddca1667d85 Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 12:45:59 +0800 Subject: [PATCH 14/50] update yml --- .github/workflows/cli.yml | 21 ++----------- .github/workflows/huigui.yml | 6 ++-- download_models.py | 59 ++++++++++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 21 deletions(-) create mode 100644 download_models.py diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 0bc08f29..7184460c 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -30,7 +30,7 @@ jobs: source activate mineru conda env list pip show coverage - # cd $GITHUB_WORKSPACE && sh tests/retry_env.sh + cd $GITHUB_WORKSPACE && sh tests/retry_env.sh cd $GITHUB_WORKSPACE && python tests/clean_coverage.py cd $GITHUB_WORKSPACE && coverage run -m pytest tests/unittest/ --cov=magic_pdf/ --cov-report html --cov-report term-missing cd $GITHUB_WORKSPACE && python tests/get_coverage.py @@ -41,22 +41,7 @@ jobs: needs: cli-test runs-on: pdf steps: - - name: get_actor - run: | - metion_list="dt-yy" - echo $GITHUB_ACTOR - if [[ $GITHUB_ACTOR == "drunkpig" ]]; then - metion_list="xuchao" - elif [[ $GITHUB_ACTOR == "myhloli" ]]; then - metion_list="zhaoxiaomeng" - elif [[ $GITHUB_ACTOR == "icecraft" ]]; then - metion_list="xurui1" - fi - echo $metion_list - echo "METIONS=$metion_list" >> "$GITHUB_ENV" - echo ${{ env.METIONS }} - - name: notify run: | - echo ${{ secrets.USER_ID }} - curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' ${{ secrets.WEBHOOK_URL }} + #echo ${{ secrets.WEBHOOK_URL }} + curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL diff --git a/.github/workflows/huigui.yml b/.github/workflows/huigui.yml index e45fa700..6e4e83a8 100644 --- a/.github/workflows/huigui.yml +++ b/.github/workflows/huigui.yml @@ -29,7 +29,7 @@ jobs: source activate mineru conda env list pip show coverage - # cd $GITHUB_WORKSPACE && sh tests/retry_env.sh + cd $GITHUB_WORKSPACE && sh tests/retry_env.sh cd $GITHUB_WORKSPACE && python tests/clean_coverage.py cd $GITHUB_WORKSPACE && coverage run -m pytest tests/unittest/ --cov=magic_pdf/ --cov-report html --cov-report term-missing cd $GITHUB_WORKSPACE && python tests/get_coverage.py @@ -57,5 +57,5 @@ jobs: - name: notify run: | - echo ${{ secrets.USER_ID }} - curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' ${{ secrets.WEBHOOK_URL }} + #echo ${{ secrets.USER_ID }} + curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL diff --git a/download_models.py b/download_models.py new file mode 100644 index 00000000..23e07608 --- /dev/null +++ b/download_models.py @@ -0,0 +1,59 @@ +import json +import os + +import requests +from modelscope import snapshot_download + + +def download_json(url): + # 下载JSON文件 + response = requests.get(url) + response.raise_for_status() # 检查请求是否成功 + return response.json() + + +def download_and_modify_json(url, local_filename, modifications): + if os.path.exists(local_filename): + data = json.load(open(local_filename)) + config_version = data.get('config_version', '0.0.0') + if config_version < '1.0.0': + data = download_json(url) + else: + data = download_json(url) + + # 修改内容 + for key, value in modifications.items(): + data[key] = value + + # 保存修改后的内容 + with open(local_filename, 'w', encoding='utf-8') as f: + json.dump(data, f, ensure_ascii=False, indent=4) + + +if __name__ == '__main__': + mineru_patterns = [ + "models/Layout/LayoutLMv3/*", + "models/Layout/YOLO/*", + "models/MFD/YOLO/*", + "models/MFR/unimernet_small/*", + "models/TabRec/TableMaster/*", + "models/TabRec/StructEqTable/*", + ] + model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns) + layoutreader_model_dir = snapshot_download('ppaanngggg/layoutreader') + model_dir = model_dir + '/models' + print(f'model_dir is: {model_dir}') + print(f'layoutreader_model_dir is: {layoutreader_model_dir}') + + json_url = 'https://gitee.com/myhloli/MinerU/raw/master/magic-pdf.template.json' + config_file_name = 'magic-pdf.json' + home_dir = os.path.expanduser('~') + config_file = os.path.join(home_dir, config_file_name) + + json_mods = { + 'models-dir': model_dir, + 'layoutreader-model-dir': layoutreader_model_dir, + } + + download_and_modify_json(json_url, config_file, json_mods) + print(f'The configuration file has been configured successfully, the path is: {config_file}') From 78e84f67641cbaf33859856c50236e799793bb1f Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 12:46:13 +0800 Subject: [PATCH 15/50] update yml --- download_models.py | 59 ---------------------------------------------- 1 file changed, 59 deletions(-) delete mode 100644 download_models.py diff --git a/download_models.py b/download_models.py deleted file mode 100644 index 23e07608..00000000 --- a/download_models.py +++ /dev/null @@ -1,59 +0,0 @@ -import json -import os - -import requests -from modelscope import snapshot_download - - -def download_json(url): - # 下载JSON文件 - response = requests.get(url) - response.raise_for_status() # 检查请求是否成功 - return response.json() - - -def download_and_modify_json(url, local_filename, modifications): - if os.path.exists(local_filename): - data = json.load(open(local_filename)) - config_version = data.get('config_version', '0.0.0') - if config_version < '1.0.0': - data = download_json(url) - else: - data = download_json(url) - - # 修改内容 - for key, value in modifications.items(): - data[key] = value - - # 保存修改后的内容 - with open(local_filename, 'w', encoding='utf-8') as f: - json.dump(data, f, ensure_ascii=False, indent=4) - - -if __name__ == '__main__': - mineru_patterns = [ - "models/Layout/LayoutLMv3/*", - "models/Layout/YOLO/*", - "models/MFD/YOLO/*", - "models/MFR/unimernet_small/*", - "models/TabRec/TableMaster/*", - "models/TabRec/StructEqTable/*", - ] - model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns) - layoutreader_model_dir = snapshot_download('ppaanngggg/layoutreader') - model_dir = model_dir + '/models' - print(f'model_dir is: {model_dir}') - print(f'layoutreader_model_dir is: {layoutreader_model_dir}') - - json_url = 'https://gitee.com/myhloli/MinerU/raw/master/magic-pdf.template.json' - config_file_name = 'magic-pdf.json' - home_dir = os.path.expanduser('~') - config_file = os.path.join(home_dir, config_file_name) - - json_mods = { - 'models-dir': model_dir, - 'layoutreader-model-dir': layoutreader_model_dir, - } - - download_and_modify_json(json_url, config_file, json_mods) - print(f'The configuration file has been configured successfully, the path is: {config_file}') From c77bec7cff497956910cad3725e192e3371d51f0 Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 12:54:16 +0800 Subject: [PATCH 16/50] update notify --- .github/workflows/cli.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 7184460c..e2f93df2 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -44,4 +44,4 @@ jobs: - name: notify run: | #echo ${{ secrets.WEBHOOK_URL }} - curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL + curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' ${{ env.WEBHOOK_URL }} From e6748482f12aecb739d41bacbd8c7245287492f4 Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 13:40:03 +0800 Subject: [PATCH 17/50] update runner env --- .github/workflows/cli.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index e2f93df2..7184460c 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -44,4 +44,4 @@ jobs: - name: notify run: | #echo ${{ secrets.WEBHOOK_URL }} - curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' ${{ env.WEBHOOK_URL }} + curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL From 8327d9d3f484a905473a0e56c9def966b0601edb Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 13:44:17 +0800 Subject: [PATCH 18/50] update runner env --- .github/workflows/cli.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 7184460c..9f01127a 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -43,5 +43,5 @@ jobs: steps: - name: notify run: | - #echo ${{ secrets.WEBHOOK_URL }} + echo $WEBHOOK_URL curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL From fc6ea7a39ac4bc7df8cf97f5ddb206c1c97db130 Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 14:07:28 +0800 Subject: [PATCH 19/50] update runner env --- .github/workflows/cli.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 9f01127a..418e5ea8 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -43,5 +43,4 @@ jobs: steps: - name: notify run: | - echo $WEBHOOK_URL curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL From cf09313b1ef19c5b8d88effc7b25dc49a8fd084e Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 14:11:11 +0800 Subject: [PATCH 20/50] update runner env --- .github/workflows/cli.yml | 2 +- .github/workflows/huigui.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/cli.yml b/.github/workflows/cli.yml index 418e5ea8..17c06480 100644 --- a/.github/workflows/cli.yml +++ b/.github/workflows/cli.yml @@ -43,4 +43,4 @@ jobs: steps: - name: notify run: | - curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL + curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'$USER_ID'"}]]}}}}' $WEBHOOK_URL diff --git a/.github/workflows/huigui.yml b/.github/workflows/huigui.yml index 6e4e83a8..8d0d55a1 100644 --- a/.github/workflows/huigui.yml +++ b/.github/workflows/huigui.yml @@ -58,4 +58,4 @@ jobs: - name: notify run: | #echo ${{ secrets.USER_ID }} - curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'${{ secrets.USER_ID }}'"}]]}}}}' $WEBHOOK_URL + curl -X POST -H "Content-Type: application/json" -d '{"msg_type":"post","content":{"post":{"zh_cn":{"title":"'${{ github.repository }}' GitHubAction Failed","content":[[{"tag":"text","text":""},{"tag":"a","text":"Please click here for details ","href":"https://github.com/'${{ github.repository }}'/actions/runs/'${GITHUB_RUN_ID}'"},{"tag":"at","user_id":"'$USER_ID'"}]]}}}}' $WEBHOOK_URL From a9723c61fe32cf9f275ec9a50b83d79c8aa9de96 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 15:07:42 +0800 Subject: [PATCH 21/50] fix(model): simplify model initialization logic --- magic_pdf/model/sub_modules/model_init.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/magic_pdf/model/sub_modules/model_init.py b/magic_pdf/model/sub_modules/model_init.py index 4d4fa1dc..b9d4d089 100644 --- a/magic_pdf/model/sub_modules/model_init.py +++ b/magic_pdf/model/sub_modules/model_init.py @@ -97,8 +97,7 @@ class AtomModelSingleton: key = (atom_model_name, layout_model_name, lang) if key not in self._models: self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) - else: - return self._models[key] + return self._models[key] def atom_model_init(model_name: str, **kwargs): From 6f636b6e7e4d3ce79c4ecb86d1a8aaf46030fe24 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 15:17:59 +0800 Subject: [PATCH 22/50] refactor(model): replace ModelSingleton with direct model initialization and improve threading - Remove usage of ModelSingleton class - Initialize model directly using custom_model_init function - Add self._lock attribute to PDFExtractKit class for thread safety- Replace local lock with self._lock for OCR processing --- magic_pdf/model/doc_analyze_by_custom_model.py | 6 ++++-- magic_pdf/model/pdf_extract_kit.py | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index a3536a3b..2bab51e8 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -143,8 +143,10 @@ def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False, if lang == "": lang = None - model_manager = ModelSingleton() - custom_model = model_manager.get_model(ocr, show_log, lang, layout_model, formula_enable, table_enable) + # model_manager = ModelSingleton() + # custom_model = model_manager.get_model(ocr, show_log, lang, layout_model, formula_enable, table_enable) + + custom_model = custom_model_init(ocr, show_log, lang, layout_model, formula_enable, table_enable) with fitz.open("pdf", pdf_bytes) as doc: pdf_page_num = doc.page_count diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 78d1930c..45371a35 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -37,6 +37,7 @@ class CustomPEKModel: """ ======== model init ======== """ + self._lock = Lock() # 获取当前文件（即 pdf_extract_kit.py）的绝对路径 current_file_path = os.path.abspath(__file__) # 获取当前文件所在的目录(model) @@ -211,14 +212,13 @@ class CustomPEKModel: # ocr识别 ocr_start = time.time() # Process each area that requires OCR processing - lock = Lock() for res in ocr_res_list: new_image, useful_list = crop_img(res, pil_img, crop_paste_x=50, crop_paste_y=50) adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(single_page_mfdetrec_res, useful_list) # OCR recognition new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR) - with lock: + with self._lock: if self.apply_ocr: ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] else: From 488660dddd6c4248a569ef902cd16aa76fd1e85f Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 15:43:41 +0800 Subject: [PATCH 23/50] refactor(model): replace AtomModelSingleton with ocr_model_init for OCR model initialization - Remove usage of AtomModelSingleton for OCR model initialization - Add import of ocr_model_init from model_init module - Update OCR model initialization process to use ocr_model_init function - Remove lock for OCR processing as it's no longer needed --- magic_pdf/model/pdf_extract_kit.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 45371a35..0970154f 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -22,7 +22,7 @@ except ImportError: from magic_pdf.config.constants import * from magic_pdf.model.model_list import AtomicModel -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init from magic_pdf.model.sub_modules.model_utils import ( clean_vram, crop_img, get_res_list_from_layout_res) from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import ( @@ -153,9 +153,14 @@ class CustomPEKModel: device=self.device, ) # 初始化ocr - self.ocr_model = atom_model_manager.get_atom_model( - atom_model_name=AtomicModel.OCR, - ocr_show_log=show_log, + # self.ocr_model = atom_model_manager.get_atom_model( + # atom_model_name=AtomicModel.OCR, + # ocr_show_log=show_log, + # det_db_box_thresh=0.3, + # lang=self.lang + # ) + self.ocr_model = ocr_model_init( + show_log=show_log, det_db_box_thresh=0.3, lang=self.lang ) @@ -218,11 +223,11 @@ class CustomPEKModel: # OCR recognition new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR) - with self._lock: - if self.apply_ocr: - ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] - else: - ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res, rec=False)[0] + # with self._lock: + if self.apply_ocr: + ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] + else: + ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res, rec=False)[0] # Integration results if ocr_res: From 30220233ab40fe1f8f36f162c9ac2327e6ae8115 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 16:11:46 +0800 Subject: [PATCH 24/50] refactor(magic_pdf): replace AtomModelSingleton with ocr_model_init for OCR model instantiation - Remove usage of AtomModelSingleton for OCR model initialization- Use ocr_model_init function for creating OCR model instance - Update import statement to include ocr_model_init- Comment out old OCR model initialization code --- magic_pdf/pdf_parse_union_core_v2.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index 2e65191e..487bc766 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -31,7 +31,7 @@ try: except ImportError: pass -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init from magic_pdf.para.para_split_v3 import para_split from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2 from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table @@ -231,10 +231,15 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang if len(empty_spans) > 0: # 初始化ocr模型 - atom_model_manager = AtomModelSingleton() - ocr_model = atom_model_manager.get_atom_model( - atom_model_name="ocr", - ocr_show_log=False, + # atom_model_manager = AtomModelSingleton() + # ocr_model = atom_model_manager.get_atom_model( + # atom_model_name="ocr", + # ocr_show_log=False, + # det_db_box_thresh=0.3, + # lang=lang + # ) + ocr_model = ocr_model_init( + show_log=False, det_db_box_thresh=0.3, lang=lang ) From 1d6000e573522cd073cd9cfd83e3ec2621a64896 Mon Sep 17 00:00:00 2001 From: dt-yy Date: Fri, 6 Dec 2024 16:29:59 +0800 Subject: [PATCH 25/50] feat: update test case --- tests/test_cli/test_cli_sdk.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/tests/test_cli/test_cli_sdk.py b/tests/test_cli/test_cli_sdk.py index d3130917..c0759f7d 100644 --- a/tests/test_cli/test_cli_sdk.py +++ b/tests/test_cli/test_cli_sdk.py @@ -7,8 +7,11 @@ from lib import common import time import magic_pdf.model as model_config from magic_pdf.pipe.UNIPipe import UNIPipe -from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter -from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter +import os +from magic_pdf.data.data_reader_writer import FileBasedDataWriter +from magic_pdf.data.data_reader_writer import S3DataReader, S3DataWriter +from magic_pdf.config.make_content_config import DropMode, MakeMode +from magic_pdf.pipe.OCRPipe import OCRPipe model_config.__use_inside_model__ = True pdf_res_path = conf.conf['pdf_res_path'] code_path = conf.conf['code_path'] @@ -41,7 +44,7 @@ class TestCli: pdf_bytes = open(pdf_path, 'rb').read() local_image_dir = os.path.join(pdf_dev_path, 'pdf', 'images') image_dir = str(os.path.basename(local_image_dir)) - image_writer = DiskReaderWriter(local_image_dir) + image_writer = FileBasedDataWriter(local_image_dir) model_json = list() jso_useful_key = {'_pdf_type': '', 'model_list': model_json} pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer) @@ -77,7 +80,7 @@ class TestCli: pdf_bytes = open(pdf_path, 'rb').read() local_image_dir = os.path.join(pdf_dev_path, 'pdf', 'images') image_dir = str(os.path.basename(local_image_dir)) - image_writer = DiskReaderWriter(local_image_dir) + image_writer = FileBasedDataWriter(local_image_dir) model_json = list() jso_useful_key = {'_pdf_type': 'ocr', 'model_list': model_json} pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer) @@ -112,7 +115,7 @@ class TestCli: pdf_bytes = open(pdf_path, 'rb').read() local_image_dir = os.path.join(pdf_dev_path, 'pdf', 'images') image_dir = str(os.path.basename(local_image_dir)) - image_writer = DiskReaderWriter(local_image_dir) + image_writer = FileBasedDataWriter(local_image_dir) model_json = list() jso_useful_key = {'_pdf_type': 'txt', 'model_list': model_json} pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer) @@ -284,12 +287,13 @@ class TestCli: pdf_endpoint = os.environ.get('pdf_endpoint', "") s3_pdf_path = conf.conf["s3_pdf_path"] image_dir = "s3://" + pdf_bucket + "/mineru/test/output" - print (image_dir) - s3pdf_cli = S3ReaderWriter(pdf_ak, pdf_sk, pdf_endpoint) - s3image_cli = S3ReaderWriter(pdf_ak, pdf_sk, pdf_endpoint, parent_path=image_dir) - pdf_bytes = s3pdf_cli.read(s3_pdf_path, mode=s3pdf_cli.MODE_BIN) - jso_useful_key = {"_pdf_type": "", "model_list": []} - pipe = UNIPipe(pdf_bytes, jso_useful_key, s3image_cli) + prefix = "mineru/test/output" + reader = S3DataReader(prefix, pdf_bucket, pdf_ak, pdf_sk, pdf_endpoint) + # = S3DataWriter(prefix, pdf_bucket, pdf_ak, pdf_sk, pdf_endpoint) + image_writer = S3DataWriter(prefix, pdf_bucket, pdf_ak, pdf_sk, pdf_endpoint) + pdf_bytes = reader.read(s3_pdf_path) + model_list = [] + pipe = OCRPipe(pdf_bytes, model_list, image_writer) pipe.pipe_classify() pipe.pipe_analyze() pipe.pipe_parse() @@ -427,3 +431,4 @@ class TestCli: if __name__ == '__main__': pytest.main() + From a1744b770f93c0fa40f71a528b37f9ac5edf2daf Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 17:10:14 +0800 Subject: [PATCH 26/50] refactor(magic_pdf): remove unused threading lock and model initialization code - Remove threading.Lock import and usage - Delete unused model initialization comments and code- Simplify OCR model initialization in both pdf_extract_kit.py and pdf_parse_union_core_v2.py --- magic_pdf/model/pdf_extract_kit.py | 11 +---------- magic_pdf/pdf_parse_union_core_v2.py | 10 ++-------- 2 files changed, 3 insertions(+), 18 deletions(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 0970154f..e9859900 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -28,8 +28,6 @@ from magic_pdf.model.sub_modules.model_utils import ( from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import ( get_adjusted_mfdetrec_res, get_ocr_result_list) -from threading import Lock - class CustomPEKModel: @@ -37,7 +35,6 @@ class CustomPEKModel: """ ======== model init ======== """ - self._lock = Lock() # 获取当前文件（即 pdf_extract_kit.py）的绝对路径 current_file_path = os.path.abspath(__file__) # 获取当前文件所在的目录(model) @@ -153,12 +150,6 @@ class CustomPEKModel: device=self.device, ) # 初始化ocr - # self.ocr_model = atom_model_manager.get_atom_model( - # atom_model_name=AtomicModel.OCR, - # ocr_show_log=show_log, - # det_db_box_thresh=0.3, - # lang=self.lang - # ) self.ocr_model = ocr_model_init( show_log=show_log, det_db_box_thresh=0.3, @@ -223,7 +214,7 @@ class CustomPEKModel: # OCR recognition new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR) - # with self._lock: + if self.apply_ocr: ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] else: diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index 487bc766..f34b295e 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -31,7 +31,7 @@ try: except ImportError: pass -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init +from magic_pdf.model.sub_modules.model_init import ocr_model_init from magic_pdf.para.para_split_v3 import para_split from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2 from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table @@ -231,13 +231,7 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang if len(empty_spans) > 0: # 初始化ocr模型 - # atom_model_manager = AtomModelSingleton() - # ocr_model = atom_model_manager.get_atom_model( - # atom_model_name="ocr", - # ocr_show_log=False, - # det_db_box_thresh=0.3, - # lang=lang - # ) + ocr_model = ocr_model_init( show_log=False, det_db_box_thresh=0.3, From 1f1335c290849707f5c94f176e6e19e7c513d382 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 17:28:24 +0800 Subject: [PATCH 27/50] build(deps): specify minimum version for ultralytics - Update `ultralytics` dependency to version >= 8.3.43 - This change ensures compatibility with yolov8 for formula detection --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index a35531f6..bf4b5e5e 100644 --- a/setup.py +++ b/setup.py @@ -39,7 +39,7 @@ if __name__ == '__main__': "full": ["unimernet==0.2.1", # unimernet升级0.2.1 "matplotlib<=3.9.0;platform_system=='Windows'", # 3.9.1及之后不提供windows的预编译包，避免一些没有编译环境的windows设备安装失败 "matplotlib;platform_system=='Linux' or platform_system=='Darwin'", # linux 和 macos 不应限制matplotlib的最高版本，以避免无法更新导致的一些bug - "ultralytics", # yolov8,公式检测 + "ultralytics>=8.3.43", # yolov8,公式检测 "paddleocr==2.7.3", # 2.8.0及2.8.1版本与detectron2有冲突，需锁定2.7.3 "paddlepaddle==3.0.0b1;platform_system=='Linux'", # 解决linux的段异常问题 "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", # windows版本3.0.0b1效率下降，需锁定2.6.1 From f2a92d5782bdc31129d9faf02fb7b342bb04c491 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 18:40:26 +0800 Subject: [PATCH 28/50] refactor(model): implement thread-safe OCR model initialization - Add threading support for OCR model initialization - Modify AtomModelSingleton to handle thread-specific instances - Update PDFExtractKit and PDFParseUnionCoreV2 to use new thread-safe OCR initialization --- magic_pdf/model/pdf_extract_kit.py | 7 ++++--- magic_pdf/model/sub_modules/model_init.py | 21 ++++++++++++++++++--- magic_pdf/pdf_parse_union_core_v2.py | 9 +++++---- 3 files changed, 27 insertions(+), 10 deletions(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index e9859900..3ab078be 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -22,7 +22,7 @@ except ImportError: from magic_pdf.config.constants import * from magic_pdf.model.model_list import AtomicModel -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton from magic_pdf.model.sub_modules.model_utils import ( clean_vram, crop_img, get_res_list_from_layout_res) from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import ( @@ -150,8 +150,9 @@ class CustomPEKModel: device=self.device, ) # 初始化ocr - self.ocr_model = ocr_model_init( - show_log=show_log, + self.ocr_model = atom_model_manager.get_atom_model( + atom_model_name=AtomicModel.OCR, + ocr_show_log=show_log, det_db_box_thresh=0.3, lang=self.lang ) diff --git a/magic_pdf/model/sub_modules/model_init.py b/magic_pdf/model/sub_modules/model_init.py index b9d4d089..31c055fb 100644 --- a/magic_pdf/model/sub_modules/model_init.py +++ b/magic_pdf/model/sub_modules/model_init.py @@ -57,6 +57,11 @@ def doclayout_yolo_model_init(weight, device='cpu'): return model +import threading +current_thread = threading.current_thread() +current_thread_id = current_thread.ident + + def ocr_model_init(show_log: bool = False, det_db_box_thresh=0.3, lang=None, @@ -92,14 +97,24 @@ class AtomModelSingleton: return cls._instance def get_atom_model(self, atom_model_name: str, **kwargs): + lang = kwargs.get('lang', None) layout_model_name = kwargs.get('layout_model_name', None) - key = (atom_model_name, layout_model_name, lang) + table_model_name = kwargs.get('table_model_name', None) + + if atom_model_name in [AtomicModel.OCR]: + key = (atom_model_name, lang, current_thread_id) + elif atom_model_name in [AtomicModel.Layout]: + key = (atom_model_name, layout_model_name) + elif atom_model_name in [AtomicModel.Table]: + key = (atom_model_name, table_model_name) + else: + key = atom_model_name + if key not in self._models: self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs) return self._models[key] - def atom_model_init(model_name: str, **kwargs): atom_model = None if model_name == AtomicModel.Layout: @@ -129,7 +144,7 @@ def atom_model_init(model_name: str, **kwargs): atom_model = ocr_model_init( kwargs.get('ocr_show_log'), kwargs.get('det_db_box_thresh'), - kwargs.get('lang') + kwargs.get('lang'), ) elif model_name == AtomicModel.Table: atom_model = table_model_init( diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index f34b295e..5146997f 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -31,7 +31,7 @@ try: except ImportError: pass -from magic_pdf.model.sub_modules.model_init import ocr_model_init +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton from magic_pdf.para.para_split_v3 import para_split from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2 from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table @@ -231,9 +231,10 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang if len(empty_spans) > 0: # 初始化ocr模型 - - ocr_model = ocr_model_init( - show_log=False, + atom_model_manager = AtomModelSingleton() + ocr_model = atom_model_manager.get_atom_model( + atom_model_name='ocr', + ocr_show_log=False, det_db_box_thresh=0.3, lang=lang ) From 47a83d28f5eeeda8973c29b22af0ecd5b00fd1ea Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 19:25:58 +0800 Subject: [PATCH 29/50] refactor(ocr): replace AtomModelSingleton with ocr_model_init for OCR model instantiation - Remove usage of AtomModelSingleton for OCR model creation - Add ocr_model_init function to initialize OCR model - Update OCR model initialization in pdf_extract_kit.py and pdf_parse_union_core_v2.py - Modify txt_spans_extract_v2 function to accept ocr_model as a parameter - Update parse_page_core function to use ocr_model instead of lang for OCR processing --- magic_pdf/model/pdf_extract_kit.py | 13 ++++++++---- magic_pdf/pdf_parse_union_core_v2.py | 31 +++++++++++++++++----------- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 3ab078be..a3bf5497 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -22,7 +22,7 @@ except ImportError: from magic_pdf.config.constants import * from magic_pdf.model.model_list import AtomicModel -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init from magic_pdf.model.sub_modules.model_utils import ( clean_vram, crop_img, get_res_list_from_layout_res) from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import ( @@ -150,9 +150,14 @@ class CustomPEKModel: device=self.device, ) # 初始化ocr - self.ocr_model = atom_model_manager.get_atom_model( - atom_model_name=AtomicModel.OCR, - ocr_show_log=show_log, + # self.ocr_model = atom_model_manager.get_atom_model( + # atom_model_name=AtomicModel.OCR, + # ocr_show_log=show_log, + # det_db_box_thresh=0.3, + # lang=self.lang + # ) + self.ocr_model = ocr_model_init( + show_log=show_log, det_db_box_thresh=0.3, lang=self.lang ) diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index 5146997f..1a0084d2 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -31,7 +31,7 @@ try: except ImportError: pass -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init from magic_pdf.para.para_split_v3 import para_split from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2 from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table @@ -152,7 +152,7 @@ def calculate_char_in_span(char_bbox, span_bbox, char, span_height_radio=0.33): return False -def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang): +def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, ocr_model): text_blocks_raw = pdf_page.get_text('rawdict', flags=fitz.TEXT_PRESERVE_WHITESPACE | fitz.TEXT_MEDIABOX_CLIP)['blocks'] @@ -231,13 +231,13 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang if len(empty_spans) > 0: # 初始化ocr模型 - atom_model_manager = AtomModelSingleton() - ocr_model = atom_model_manager.get_atom_model( - atom_model_name='ocr', - ocr_show_log=False, - det_db_box_thresh=0.3, - lang=lang - ) + # atom_model_manager = AtomModelSingleton() + # ocr_model = atom_model_manager.get_atom_model( + # atom_model_name='ocr', + # ocr_show_log=False, + # det_db_box_thresh=0.3, + # lang=lang + # ) for span in empty_spans: # 对span的bbox截图再ocr @@ -613,7 +613,7 @@ def remove_outside_spans(spans, all_bboxes, all_discarded_blocks): def parse_page_core( - page_doc: PageableData, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, lang + page_doc: PageableData, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, ocr_model ): need_drop = False drop_reason = [] @@ -682,7 +682,7 @@ def parse_page_core( if parse_mode == SupportedPdfParseMethod.TXT: """使用新版本的混合ocr方案""" - spans = txt_spans_extract_v2(page_doc, spans, all_bboxes, all_discarded_blocks, lang) + spans = txt_spans_extract_v2(page_doc, spans, all_bboxes, all_discarded_blocks, ocr_model) elif parse_mode == SupportedPdfParseMethod.OCR: pass @@ -771,6 +771,13 @@ def pdf_parse_union( debug_mode=False, lang=None, ): + + ocr_model = ocr_model_init( + show_log=False, + det_db_box_thresh=0.3, + lang=lang + ) + pdf_bytes_md5 = compute_md5(dataset.data_bits()) """初始化空的pdf_info_dict""" @@ -806,7 +813,7 @@ def pdf_parse_union( """解析pdf中的每一页""" if start_page_id <= page_id <= end_page_id: page_info = parse_page_core( - page, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, lang + page, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, ocr_model ) else: page_info = page.get_page_info() From 012a46e07d51ce0e55d929fdc823e9c046db1489 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 6 Dec 2024 20:35:43 +0800 Subject: [PATCH 30/50] refactor(magic-pdf): optimize model initialization and concurrency control - Remove concurrency limit logic from app.py - Update model initialization process in various modules - Remove unused VRAM check for concurrency limit - Refactor OCR model initialization in pdf_extract_kit.py - Update txt_spans_extract_v2 function to use lang parameter instead of ocr_model --- .../model/doc_analyze_by_custom_model.py | 6 ++-- magic_pdf/model/pdf_extract_kit.py | 13 +++------ magic_pdf/model/sub_modules/model_init.py | 7 +---- magic_pdf/pdf_parse_union_core_v2.py | 28 ++++++++----------- projects/gradio_app/app.py | 14 +--------- 5 files changed, 19 insertions(+), 49 deletions(-) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index 2bab51e8..a3536a3b 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -143,10 +143,8 @@ def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False, if lang == "": lang = None - # model_manager = ModelSingleton() - # custom_model = model_manager.get_model(ocr, show_log, lang, layout_model, formula_enable, table_enable) - - custom_model = custom_model_init(ocr, show_log, lang, layout_model, formula_enable, table_enable) + model_manager = ModelSingleton() + custom_model = model_manager.get_model(ocr, show_log, lang, layout_model, formula_enable, table_enable) with fitz.open("pdf", pdf_bytes) as doc: pdf_page_num = doc.page_count diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index a3bf5497..3ab078be 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -22,7 +22,7 @@ except ImportError: from magic_pdf.config.constants import * from magic_pdf.model.model_list import AtomicModel -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton from magic_pdf.model.sub_modules.model_utils import ( clean_vram, crop_img, get_res_list_from_layout_res) from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import ( @@ -150,14 +150,9 @@ class CustomPEKModel: device=self.device, ) # 初始化ocr - # self.ocr_model = atom_model_manager.get_atom_model( - # atom_model_name=AtomicModel.OCR, - # ocr_show_log=show_log, - # det_db_box_thresh=0.3, - # lang=self.lang - # ) - self.ocr_model = ocr_model_init( - show_log=show_log, + self.ocr_model = atom_model_manager.get_atom_model( + atom_model_name=AtomicModel.OCR, + ocr_show_log=show_log, det_db_box_thresh=0.3, lang=self.lang ) diff --git a/magic_pdf/model/sub_modules/model_init.py b/magic_pdf/model/sub_modules/model_init.py index 31c055fb..81571bc3 100644 --- a/magic_pdf/model/sub_modules/model_init.py +++ b/magic_pdf/model/sub_modules/model_init.py @@ -57,11 +57,6 @@ def doclayout_yolo_model_init(weight, device='cpu'): return model -import threading -current_thread = threading.current_thread() -current_thread_id = current_thread.ident - - def ocr_model_init(show_log: bool = False, det_db_box_thresh=0.3, lang=None, @@ -103,7 +98,7 @@ class AtomModelSingleton: table_model_name = kwargs.get('table_model_name', None) if atom_model_name in [AtomicModel.OCR]: - key = (atom_model_name, lang, current_thread_id) + key = (atom_model_name, lang) elif atom_model_name in [AtomicModel.Layout]: key = (atom_model_name, layout_model_name) elif atom_model_name in [AtomicModel.Table]: diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index 1a0084d2..ccaf9c1c 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -152,7 +152,7 @@ def calculate_char_in_span(char_bbox, span_bbox, char, span_height_radio=0.33): return False -def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, ocr_model): +def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang): text_blocks_raw = pdf_page.get_text('rawdict', flags=fitz.TEXT_PRESERVE_WHITESPACE | fitz.TEXT_MEDIABOX_CLIP)['blocks'] @@ -231,13 +231,13 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, ocr_ if len(empty_spans) > 0: # 初始化ocr模型 - # atom_model_manager = AtomModelSingleton() - # ocr_model = atom_model_manager.get_atom_model( - # atom_model_name='ocr', - # ocr_show_log=False, - # det_db_box_thresh=0.3, - # lang=lang - # ) + atom_model_manager = AtomModelSingleton() + ocr_model = atom_model_manager.get_atom_model( + atom_model_name='ocr', + ocr_show_log=False, + det_db_box_thresh=0.3, + lang=lang + ) for span in empty_spans: # 对span的bbox截图再ocr @@ -613,7 +613,7 @@ def remove_outside_spans(spans, all_bboxes, all_discarded_blocks): def parse_page_core( - page_doc: PageableData, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, ocr_model + page_doc: PageableData, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, lang ): need_drop = False drop_reason = [] @@ -682,7 +682,7 @@ def parse_page_core( if parse_mode == SupportedPdfParseMethod.TXT: """使用新版本的混合ocr方案""" - spans = txt_spans_extract_v2(page_doc, spans, all_bboxes, all_discarded_blocks, ocr_model) + spans = txt_spans_extract_v2(page_doc, spans, all_bboxes, all_discarded_blocks, lang) elif parse_mode == SupportedPdfParseMethod.OCR: pass @@ -772,12 +772,6 @@ def pdf_parse_union( lang=None, ): - ocr_model = ocr_model_init( - show_log=False, - det_db_box_thresh=0.3, - lang=lang - ) - pdf_bytes_md5 = compute_md5(dataset.data_bits()) """初始化空的pdf_info_dict""" @@ -813,7 +807,7 @@ def pdf_parse_union( """解析pdf中的每一页""" if start_page_id <= page_id <= end_page_id: page_info = parse_page_core( - page, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, ocr_model + page, magic_model, page_id, pdf_bytes_md5, imageWriter, parse_mode, lang ) else: page_info = page.get_page_info() diff --git a/projects/gradio_app/app.py b/projects/gradio_app/app.py index 1dffdafb..1ccede3b 100644 --- a/projects/gradio_app/app.py +++ b/projects/gradio_app/app.py @@ -14,9 +14,7 @@ from gradio_pdf import PDF from loguru import logger from magic_pdf.data.data_reader_writer import FileBasedDataReader -from magic_pdf.libs.config_reader import get_device from magic_pdf.libs.hash_utils import compute_sha256 -from magic_pdf.model.sub_modules.model_utils import get_vram from magic_pdf.tools.common import do_parse, prepare_env @@ -185,16 +183,6 @@ def to_pdf(file_path): return tmp_file_path -def get_concurrency_limit(vram_threshold=7.5): - vram = get_vram(device = get_device()) - if vram is not None and isinstance(vram, (int, float)): - concurrency_limit = max(1, int(vram // vram_threshold)) - else: - concurrency_limit = 1 - # logger.info(f'concurrency_limit: {concurrency_limit}') - return concurrency_limit - - if __name__ == '__main__': with gr.Blocks() as demo: gr.HTML(header) @@ -231,7 +219,7 @@ if __name__ == '__main__': md_text = gr.TextArea(lines=45, show_copy_button=True) file.upload(fn=to_pdf, inputs=file, outputs=pdf_show) change_bu.click(fn=to_markdown, inputs=[pdf_show, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language], - outputs=[md, md_text, output_file, pdf_show], concurrency_limit=get_concurrency_limit()) + outputs=[md, md_text, output_file, pdf_show]) clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr, table_enable, language]) demo.launch(server_name='0.0.0.0') From 74ee428bbb45638dd011a796c26ecbae8954e7a1 Mon Sep 17 00:00:00 2001 From: myhloli Date: Sat, 7 Dec 2024 17:53:54 +0800 Subject: [PATCH 31/50] fix(dict2md): add space for inline equations in CJK contexts - In Chinese, Japanese, and Korean (CJK) languages, no space is needed for line breaks within paragraphs. - However, if an inline equation is at the end of a line, a space should be added to separate it from the following text. - This change improves the formatting of documents containing both CJK text and inline equations. --- magic_pdf/dict2md/ocr_mkcontent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/magic_pdf/dict2md/ocr_mkcontent.py b/magic_pdf/dict2md/ocr_mkcontent.py index 9c83fdb0..ad7866b2 100644 --- a/magic_pdf/dict2md/ocr_mkcontent.py +++ b/magic_pdf/dict2md/ocr_mkcontent.py @@ -165,8 +165,8 @@ def merge_para_with_text(para_block): if content: langs = ['zh', 'ja', 'ko'] # logger.info(f'block_lang: {block_lang}, content: {content}') - if block_lang in langs: # 中文/日语/韩文语境下，换行不需要空格分隔 - if j == len(line['spans']) - 1: + if block_lang in langs: # 中文/日语/韩文语境下，换行不需要空格分隔,但是如果是行内公式结尾，还是要加空格 + if j == len(line['spans']) - 1 and span_type not in [ContentType.InlineEquation]: para_text += content else: para_text += f'{content} ' From 87af738ab1364c6751644a0b2efbd5b20364c06b Mon Sep 17 00:00:00 2001 From: sawmice Date: Sat, 7 Dec 2024 18:03:39 +0800 Subject: [PATCH 32/50] fix: 1. ocr txt mode error 2. lose pdf_parse_type field --- magic_pdf/config/constants.py | 5 +++++ magic_pdf/model/operators.py | 17 ++++++++++++++--- magic_pdf/user_api.py | 4 +--- 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/magic_pdf/config/constants.py b/magic_pdf/config/constants.py index 8a0661d9..c8a63bd1 100644 --- a/magic_pdf/config/constants.py +++ b/magic_pdf/config/constants.py @@ -51,3 +51,8 @@ class MODEL_NAME: UniMerNet_v2_Small = 'unimernet_small' RAPID_TABLE = 'rapid_table' + + +PARSE_TYPE_TXT = 'txt' +PARSE_TYPE_OCR = 'ocr' + diff --git a/magic_pdf/model/operators.py b/magic_pdf/model/operators.py index a69cf981..c410353a 100644 --- a/magic_pdf/model/operators.py +++ b/magic_pdf/model/operators.py @@ -6,12 +6,14 @@ from typing import Callable from magic_pdf.config.enums import SupportedPdfParseMethod from magic_pdf.data.data_reader_writer import DataWriter from magic_pdf.data.dataset import Dataset +from magic_pdf.libs.version import __version__ from magic_pdf.filter import classify from magic_pdf.libs.draw_bbox import draw_model_bbox from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union from magic_pdf.pipe.operators import PipeResult from magic_pdf.model import InferenceResultBase - +from magic_pdf.libs.version import __version__ +from magic_pdf.config.constants import PARSE_TYPE_TXT, PARSE_TYPE_OCR class InferenceResult(InferenceResultBase): def __init__(self, inference_results: list, dataset: Dataset): @@ -129,7 +131,7 @@ class InferenceResult(InferenceResultBase): res = pdf_parse_union(*args, **kwargs) return PipeResult(res, self._dataset) - return self.apply( + res = self.apply( proc, self._dataset, imageWriter, @@ -139,6 +141,11 @@ class InferenceResult(InferenceResultBase): debug_mode=debug_mode, lang=lang, ) + res['_parse_type'] = PARSE_TYPE_TXT + res['_version_name'] = __version__ + + return res + def pipe_ocr_mode( self, @@ -166,7 +173,7 @@ class InferenceResult(InferenceResultBase): res = pdf_parse_union(*args, **kwargs) return PipeResult(res, self._dataset) - return self.apply( + res = self.apply( proc, self._dataset, imageWriter, @@ -176,3 +183,7 @@ class InferenceResult(InferenceResultBase): debug_mode=debug_mode, lang=lang, ) + res['_parse_type'] = PARSE_TYPE_OCR + + res['_version_name'] = __version__ + return res \ No newline at end of file diff --git a/magic_pdf/user_api.py b/magic_pdf/user_api.py index 99b9502f..156bf649 100644 --- a/magic_pdf/user_api.py +++ b/magic_pdf/user_api.py @@ -15,9 +15,7 @@ from magic_pdf.libs.version import __version__ from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze 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' +from magic_pdf.config.constants import PARSE_TYPE_TXT, PARSE_TYPE_OCR def parse_txt_pdf( From 57f9f9dcf94777ad0ac26c78287c10edcb1ff264 Mon Sep 17 00:00:00 2001 From: icecraft Date: Mon, 9 Dec 2024 11:20:21 +0800 Subject: [PATCH 33/50] fix: add parse_pdf_type and version --- magic_pdf/model/operators.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/magic_pdf/model/operators.py b/magic_pdf/model/operators.py index c410353a..189fa8a5 100644 --- a/magic_pdf/model/operators.py +++ b/magic_pdf/model/operators.py @@ -3,17 +3,17 @@ import json import os from typing import Callable +from magic_pdf.config.constants import PARSE_TYPE_OCR, PARSE_TYPE_TXT from magic_pdf.config.enums import SupportedPdfParseMethod from magic_pdf.data.data_reader_writer import DataWriter from magic_pdf.data.dataset import Dataset -from magic_pdf.libs.version import __version__ from magic_pdf.filter import classify from magic_pdf.libs.draw_bbox import draw_model_bbox +from magic_pdf.libs.version import __version__ +from magic_pdf.model import InferenceResultBase from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union from magic_pdf.pipe.operators import PipeResult -from magic_pdf.model import InferenceResultBase -from magic_pdf.libs.version import __version__ -from magic_pdf.config.constants import PARSE_TYPE_TXT, PARSE_TYPE_OCR + class InferenceResult(InferenceResultBase): def __init__(self, inference_results: list, dataset: Dataset): @@ -129,6 +129,10 @@ class InferenceResult(InferenceResultBase): def proc(*args, **kwargs) -> PipeResult: res = pdf_parse_union(*args, **kwargs) + res['_parse_type'] = PARSE_TYPE_TXT + res['_version_name'] = __version__ + if 'lang' in kwargs and kwargs['lang'] is not None: + res['lang'] = kwargs['lang'] return PipeResult(res, self._dataset) res = self.apply( @@ -141,11 +145,7 @@ class InferenceResult(InferenceResultBase): debug_mode=debug_mode, lang=lang, ) - res['_parse_type'] = PARSE_TYPE_TXT - res['_version_name'] = __version__ - return res - def pipe_ocr_mode( self, @@ -171,19 +171,20 @@ class InferenceResult(InferenceResultBase): def proc(*args, **kwargs) -> PipeResult: res = pdf_parse_union(*args, **kwargs) + res['_parse_type'] = PARSE_TYPE_OCR + res['_version_name'] = __version__ + if 'lang' in kwargs and kwargs['lang'] is not None: + res['lang'] = kwargs['lang'] return PipeResult(res, self._dataset) res = self.apply( proc, self._dataset, imageWriter, - SupportedPdfParseMethod.TXT, + SupportedPdfParseMethod.OCR, start_page_id=start_page_id, end_page_id=end_page_id, debug_mode=debug_mode, lang=lang, ) - res['_parse_type'] = PARSE_TYPE_OCR - - res['_version_name'] = __version__ - return res \ No newline at end of file + return res From 113448903aac099e42522dc4e2e2d80d0f0dc45c Mon Sep 17 00:00:00 2001 From: icecraft Date: Mon, 9 Dec 2024 13:58:27 +0800 Subject: [PATCH 34/50] fix: unicode decode error --- magic_pdf/data/data_reader_writer/base.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/magic_pdf/data/data_reader_writer/base.py b/magic_pdf/data/data_reader_writer/base.py index 7c9a8e8e..d294b329 100644 --- a/magic_pdf/data/data_reader_writer/base.py +++ b/magic_pdf/data/data_reader_writer/base.py @@ -48,4 +48,16 @@ class DataWriter(ABC): path (str): the target file where to write data (str): the data want to write """ - self.write(path, data.encode()) + + def safe_encode(data: str, method: str): + try: + bit_data = data.encode(encoding=method, errors='replace') + return bit_data, True + except: # noqa + return None, False + + for method in ['utf-8', 'ascii']: + bit_data, flag = safe_encode(data, method) + if flag: + self.write(path, bit_data) + break From 2ae103940885f2337dfb35e15994df968473ea85 Mon Sep 17 00:00:00 2001 From: myhloli Date: Mon, 9 Dec 2024 14:26:48 +0800 Subject: [PATCH 35/50] build(deps): update dependency versions - Update ultralytics to >=8.3.47 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index bf4b5e5e..e9484bb9 100644 --- a/setup.py +++ b/setup.py @@ -39,7 +39,7 @@ if __name__ == '__main__': "full": ["unimernet==0.2.1", # unimernet升级0.2.1 "matplotlib<=3.9.0;platform_system=='Windows'", # 3.9.1及之后不提供windows的预编译包，避免一些没有编译环境的windows设备安装失败 "matplotlib;platform_system=='Linux' or platform_system=='Darwin'", # linux 和 macos 不应限制matplotlib的最高版本，以避免无法更新导致的一些bug - "ultralytics>=8.3.43", # yolov8,公式检测 + "ultralytics>=8.3.47", # yolov8,公式检测 "paddleocr==2.7.3", # 2.8.0及2.8.1版本与detectron2有冲突，需锁定2.7.3 "paddlepaddle==3.0.0b1;platform_system=='Linux'", # 解决linux的段异常问题 "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", # windows版本3.0.0b1效率下降，需锁定2.6.1 From a296ea41f9d14cda6257eca51ce089531029a2ba Mon Sep 17 00:00:00 2001 From: myhloli Date: Mon, 9 Dec 2024 18:08:27 +0800 Subject: [PATCH 36/50] refactor(magic_pdf): optimize environment setup and dependencies - Add environment variables to disable albumentations and yolo updates - Import torchtext and disable deprecation warnings - Update unimernet to 0.2.2 - Specify ultralytics version as >=8.3.48 - Remove upper version limit for torch --- magic_pdf/model/doc_analyze_by_custom_model.py | 13 ++++++++++++- requirements-docker.txt | 4 ++-- requirements.txt | 2 +- setup.py | 4 ++-- 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index 17a314ff..38605357 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -1,10 +1,21 @@ - +import os import time import fitz import numpy as np from loguru import logger +os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新 +os.environ['YOLO_VERBOSE'] = 'False' # disable yolo logger + +try: + import torchtext + + if torchtext.__version__ >= '0.18.0': + torchtext.disable_torchtext_deprecation_warning() +except ImportError: + pass + import magic_pdf.model as model_config from magic_pdf.data.dataset import Dataset from magic_pdf.libs.clean_memory import clean_memory diff --git a/requirements-docker.txt b/requirements-docker.txt index 34fd61e2..74269994 100644 --- a/requirements-docker.txt +++ b/requirements-docker.txt @@ -7,9 +7,9 @@ numpy>=1.21.6,<2.0.0 fast-langdetect==0.2.0 scikit-learn>=1.0.2 pdfminer.six==20231228 -unimernet==0.2.1 +unimernet==0.2.2 matplotlib -ultralytics +ultralytics>=8.3.48 paddleocr==2.7.3 paddlepaddle==3.0.0b1 struct-eqtable==0.3.2 diff --git a/requirements.txt b/requirements.txt index d308e30a..c0e511a5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,7 @@ numpy>=1.21.6,<2.0.0 pydantic>=2.7.2,<2.8.0 PyMuPDF>=1.24.9 scikit-learn>=1.0.2 -torch>=2.2.2,<=2.3.1 +torch>=2.2.2 transformers # pdfminer.six==20231228 # The requirements.txt must ensure that only necessary external dependencies are introduced. If there are new dependencies to add, please contact the project administrator. diff --git a/setup.py b/setup.py index e9484bb9..a1f84a0e 100644 --- a/setup.py +++ b/setup.py @@ -36,10 +36,10 @@ if __name__ == '__main__': "paddlepaddle==3.0.0b1;platform_system=='Linux'", "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", ], - "full": ["unimernet==0.2.1", # unimernet升级0.2.1 + "full": ["unimernet==0.2.2", # unimernet升级0.2.1 "matplotlib<=3.9.0;platform_system=='Windows'", # 3.9.1及之后不提供windows的预编译包，避免一些没有编译环境的windows设备安装失败 "matplotlib;platform_system=='Linux' or platform_system=='Darwin'", # linux 和 macos 不应限制matplotlib的最高版本，以避免无法更新导致的一些bug - "ultralytics>=8.3.47", # yolov8,公式检测 + "ultralytics>=8.3.48", # yolov8,公式检测 "paddleocr==2.7.3", # 2.8.0及2.8.1版本与detectron2有冲突，需锁定2.7.3 "paddlepaddle==3.0.0b1;platform_system=='Linux'", # 解决linux的段异常问题 "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", # windows版本3.0.0b1效率下降，需锁定2.6.1 From ede7d361a4f085db7d023c1cf5c121b011d16ef5 Mon Sep 17 00:00:00 2001 From: myhloli Date: Mon, 9 Dec 2024 18:11:24 +0800 Subject: [PATCH 37/50] docs(windows): update CUDA installation guide - Remove specific version requirements for torch and torchvision - Simplify installation command in both English and Chinese guides - Delete important note about version compatibility --- docs/README_Windows_CUDA_Acceleration_en_US.md | 10 +--------- docs/README_Windows_CUDA_Acceleration_zh_CN.md | 11 +---------- 2 files changed, 2 insertions(+), 19 deletions(-) diff --git a/docs/README_Windows_CUDA_Acceleration_en_US.md b/docs/README_Windows_CUDA_Acceleration_en_US.md index 73a2de5c..e8252145 100644 --- a/docs/README_Windows_CUDA_Acceleration_en_US.md +++ b/docs/README_Windows_CUDA_Acceleration_en_US.md @@ -65,16 +65,8 @@ If your graphics card has at least 8GB of VRAM, follow these steps to test CUDA- 1. **Overwrite the installation of torch and torchvision** supporting CUDA. ``` - pip install --force-reinstall torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu118 + pip install --force-reinstall torch torchvision --index-url https://download.pytorch.org/whl/cu118 ``` - > [!IMPORTANT] - > Ensure the following versions are specified in the command: - > - > ``` - > torch==2.3.1 torchvision==0.18.1 - > ``` - > - > These are the highest versions we support. Installing higher versions without specifying them will cause the program to fail. 2. **Modify the value of `"device-mode"`** in the `magic-pdf.json` configuration file located in your user directory. diff --git a/docs/README_Windows_CUDA_Acceleration_zh_CN.md b/docs/README_Windows_CUDA_Acceleration_zh_CN.md index d0454f7d..0b156994 100644 --- a/docs/README_Windows_CUDA_Acceleration_zh_CN.md +++ b/docs/README_Windows_CUDA_Acceleration_zh_CN.md @@ -66,18 +66,9 @@ pip install -U magic-pdf[full] --extra-index-url https://wheels.myhloli.com -i h **1.覆盖安装支持cuda的torch和torchvision** ```bash -pip install --force-reinstall torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu118 +pip install --force-reinstall torch torchvision --index-url https://download.pytorch.org/whl/cu118 ``` -> [!IMPORTANT] -> 务必在命令中指定以下版本 -> -> ```bash -> torch==2.3.1 torchvision==0.18.1 -> ``` -> -> 这是我们支持的最高版本，如果不指定版本会自动安装更高版本导致程序无法运行 - **2.修改【用户目录】中配置文件magic-pdf.json中"device-mode"的值** ```json From 2b6e94424cbc8f4d52972305d6a8c1cf6c45e3a3 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 10 Dec 2024 17:50:51 +0800 Subject: [PATCH 38/50] refactor: comment out clean_memory function call - Remove the call to clean_memory() function from pdf_parse_union_core_v2.py - This change may affect memory usage and needs to be tested to ensure proper functionality --- magic_pdf/pdf_parse_union_core_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index b226e8d3..576fc2b3 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -825,7 +825,7 @@ def pdf_parse_union( 'pdf_info': pdf_info_list, } - clean_memory() + # clean_memory() return new_pdf_info_dict From dd7f6781f6236e0ab7acad668d412ef27c050287 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 10 Dec 2024 18:11:59 +0800 Subject: [PATCH 39/50] fix(magic_pdf): disable PaddlePaddle signal handler - Import paddle module and disable its signal handler to prevent interference with other components - This change addresses potential conflicts between PaddlePaddle and other libraries or system signals --- magic_pdf/model/doc_analyze_by_custom_model.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/magic_pdf/model/doc_analyze_by_custom_model.py b/magic_pdf/model/doc_analyze_by_custom_model.py index 38605357..ec46d867 100644 --- a/magic_pdf/model/doc_analyze_by_custom_model.py +++ b/magic_pdf/model/doc_analyze_by_custom_model.py @@ -5,6 +5,10 @@ import fitz import numpy as np from loguru import logger +# 关闭paddle的信号处理 +import paddle +paddle.disable_signal_handler() + os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新 os.environ['YOLO_VERBOSE'] = 'False' # disable yolo logger From 52dfdd535417fe22ca190d1c416ba033b2e0ec55 Mon Sep 17 00:00:00 2001 From: myhloli Date: Tue, 10 Dec 2024 20:05:22 +0800 Subject: [PATCH 40/50] refactor(magic_pdf): remove unnecessary comment - Remove commented-out call to clean_memory() function - This change simplifies the code by eliminating an unused code snippet --- magic_pdf/pdf_parse_union_core_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index 576fc2b3..b226e8d3 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -825,7 +825,7 @@ def pdf_parse_union( 'pdf_info': pdf_info_list, } - # clean_memory() + clean_memory() return new_pdf_info_dict From 56fad23d67a654a2b033e36638ee35095e5c8a51 Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 01:24:02 +0800 Subject: [PATCH 41/50] refactor(magic_pdf): update paddleocr module import paths - Modify import paths for paddleocr utilities in ocr_utils.py and ppocr_273_mod.py - Change from `ppocr.utils.utility` to `paddleocr.ppocr.utils.utility` - Update related import statements in two files to reflect the new path --- magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py | 2 +- .../model/sub_modules/ocr/paddleocr/ppocr_273_mod.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py b/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py index 11526d53..75ce8773 100644 --- a/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py +++ b/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py @@ -7,7 +7,7 @@ import base64 from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold from magic_pdf.pre_proc.ocr_dict_merge import merge_spans_to_line -from ppocr.utils.utility import check_and_read +from paddleocr.ppocr.utils.utility import check_and_read def img_decode(content: bytes): diff --git a/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py b/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py index d900254b..70ffc3e9 100644 --- a/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py +++ b/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py @@ -4,10 +4,10 @@ import cv2 import numpy as np from paddleocr import PaddleOCR -from ppocr.utils.logging import get_logger -from ppocr.utils.utility import alpha_to_color, binarize_img -from tools.infer.predict_system import sorted_boxes -from tools.infer.utility import get_rotate_crop_image, get_minarea_rect_crop +from paddleocr.ppocr.utils.logging import get_logger +from paddleocr.ppocr.utils.utility import alpha_to_color, binarize_img +from paddleocr.tools.infer.predict_system import sorted_boxes +from paddleocr.tools.infer.utility import get_rotate_crop_image, get_minarea_rect_crop from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import update_det_boxes, merge_det_boxes, check_img From 01cd633ded9a95215c8efdc0f680330971258efb Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 01:26:18 +0800 Subject: [PATCH 42/50] refactor(tablemaster): update import paths for TableSystem and init_args - Change import path for TableSystem from 'ppstructure.table.predict_table' to 'paddleocr.ppstructure.table.predict_table' - Change import path for init_args from 'ppstructure.utility' to 'paddleocr.ppstructure.utility' --- .../model/sub_modules/table/tablemaster/tablemaster_paddle.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py b/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py index 88dfbaf7..efadbdd7 100644 --- a/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py +++ b/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py @@ -2,8 +2,8 @@ import os import cv2 import numpy as np -from ppstructure.table.predict_table import TableSystem -from ppstructure.utility import init_args +from paddleocr.ppstructure.table.predict_table import TableSystem +from paddleocr.ppstructure.utility import init_args from PIL import Image from magic_pdf.config.constants import * # noqa: F403 From e1be7da644afdce24e67ae419e281d8297b4aa1f Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 01:41:35 +0800 Subject: [PATCH 43/50] refactor(magic_pdf): switch to pdfminer for invalid character detection - Replace MuPDF with pdfminer for detecting invalid characters in PDFs - Uncomment and update the detect_invalid_chars function to use pdfminer - Update the check_invalid_chars function in pdf_meta_scan.py to use the new implementation --- magic_pdf/filter/pdf_meta_scan.py | 5 +-- magic_pdf/libs/pdf_check.py | 60 +++++++++++++++---------------- 2 files changed, 33 insertions(+), 32 deletions(-) diff --git a/magic_pdf/filter/pdf_meta_scan.py b/magic_pdf/filter/pdf_meta_scan.py index 400c1707..eecad515 100644 --- a/magic_pdf/filter/pdf_meta_scan.py +++ b/magic_pdf/filter/pdf_meta_scan.py @@ -8,7 +8,7 @@ from loguru import logger from magic_pdf.config.drop_reason import DropReason from magic_pdf.libs.commons import get_top_percent_list, mymax from magic_pdf.libs.language import detect_lang -from magic_pdf.libs.pdf_check import detect_invalid_chars_by_pymupdf +from magic_pdf.libs.pdf_check import detect_invalid_chars_by_pymupdf, detect_invalid_chars scan_max_page = 50 junk_limit_min = 10 @@ -323,7 +323,8 @@ def get_language(doc: fitz.Document): def check_invalid_chars(pdf_bytes): """乱码检测.""" - return detect_invalid_chars_by_pymupdf(pdf_bytes) + # return detect_invalid_chars_by_pymupdf(pdf_bytes) + return detect_invalid_chars(pdf_bytes) def pdf_meta_scan(pdf_bytes: bytes): diff --git a/magic_pdf/libs/pdf_check.py b/magic_pdf/libs/pdf_check.py index fe40a45a..d4d3c24b 100644 --- a/magic_pdf/libs/pdf_check.py +++ b/magic_pdf/libs/pdf_check.py @@ -1,9 +1,9 @@ import fitz import numpy as np from loguru import logger -# import re -# from io import BytesIO -# from pdfminer.high_level import extract_text +import re +from io import BytesIO +from pdfminer.high_level import extract_text def calculate_sample_count(total_page: int): @@ -33,33 +33,33 @@ def extract_pages(src_pdf_bytes: bytes) -> fitz.Document: return sample_docs -# def detect_invalid_chars(src_pdf_bytes: bytes) -> bool: -# """" -# 检测PDF中是否包含非法字符 -# """ -# '''pdfminer比较慢,需要先随机抽取10页左右的sample''' -# sample_docs = extract_pages(src_pdf_bytes) -# sample_pdf_bytes = sample_docs.tobytes() -# sample_pdf_file_like_object = BytesIO(sample_pdf_bytes) -# text = extract_text(sample_pdf_file_like_object) -# text = text.replace("\n", "") -# # logger.info(text) -# '''乱码文本用pdfminer提取出来的文本特征是(cid:xxx)''' -# cid_pattern = re.compile(r'\(cid:\d+\)') -# matches = cid_pattern.findall(text) -# cid_count = len(matches) -# cid_len = sum(len(match) for match in matches) -# text_len = len(text) -# if text_len == 0: -# cid_chars_radio = 0 -# else: -# cid_chars_radio = cid_count/(cid_count + text_len - cid_len) -# logger.info(f"cid_count: {cid_count}, text_len: {text_len}, cid_chars_radio: {cid_chars_radio}") -# '''当一篇文章存在5%以上的文本是乱码时,认为该文档为乱码文档''' -# if cid_chars_radio > 0.05: -# return False # 乱码文档 -# else: -# return True # 正常文档 +def detect_invalid_chars(src_pdf_bytes: bytes) -> bool: + """" + 检测PDF中是否包含非法字符 + """ + '''pdfminer比较慢,需要先随机抽取10页左右的sample''' + sample_docs = extract_pages(src_pdf_bytes) + sample_pdf_bytes = sample_docs.tobytes() + sample_pdf_file_like_object = BytesIO(sample_pdf_bytes) + text = extract_text(sample_pdf_file_like_object) + text = text.replace("\n", "") + # logger.info(text) + '''乱码文本用pdfminer提取出来的文本特征是(cid:xxx)''' + cid_pattern = re.compile(r'\(cid:\d+\)') + matches = cid_pattern.findall(text) + cid_count = len(matches) + cid_len = sum(len(match) for match in matches) + text_len = len(text) + if text_len == 0: + cid_chars_radio = 0 + else: + cid_chars_radio = cid_count/(cid_count + text_len - cid_len) + logger.info(f"cid_count: {cid_count}, text_len: {text_len}, cid_chars_radio: {cid_chars_radio}") + '''当一篇文章存在5%以上的文本是乱码时,认为该文档为乱码文档''' + if cid_chars_radio > 0.05: + return False # 乱码文档 + else: + return True # 正常文档 def count_replacement_characters(text: str) -> int: From 061c03a0db8bab6e0aeca40ce8584598f62b50fb Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 01:49:24 +0800 Subject: [PATCH 44/50] refactor(model): update import paths for PaddleOCR modules - Change import paths from paddleocr.ppocr to ppocr for utility functions - Update import paths for logging and utility modules in ppocr_273_mod.py- Modify import paths for tablemaster_paddle.py to use ppstructure instead of paddleocr.ppstructure --- magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py | 2 +- .../model/sub_modules/ocr/paddleocr/ppocr_273_mod.py | 8 ++++---- .../sub_modules/table/tablemaster/tablemaster_paddle.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py b/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py index 75ce8773..11526d53 100644 --- a/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py +++ b/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py @@ -7,7 +7,7 @@ import base64 from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold from magic_pdf.pre_proc.ocr_dict_merge import merge_spans_to_line -from paddleocr.ppocr.utils.utility import check_and_read +from ppocr.utils.utility import check_and_read def img_decode(content: bytes): diff --git a/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py b/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py index 70ffc3e9..d900254b 100644 --- a/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py +++ b/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py @@ -4,10 +4,10 @@ import cv2 import numpy as np from paddleocr import PaddleOCR -from paddleocr.ppocr.utils.logging import get_logger -from paddleocr.ppocr.utils.utility import alpha_to_color, binarize_img -from paddleocr.tools.infer.predict_system import sorted_boxes -from paddleocr.tools.infer.utility import get_rotate_crop_image, get_minarea_rect_crop +from ppocr.utils.logging import get_logger +from ppocr.utils.utility import alpha_to_color, binarize_img +from tools.infer.predict_system import sorted_boxes +from tools.infer.utility import get_rotate_crop_image, get_minarea_rect_crop from magic_pdf.model.sub_modules.ocr.paddleocr.ocr_utils import update_det_boxes, merge_det_boxes, check_img diff --git a/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py b/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py index efadbdd7..88dfbaf7 100644 --- a/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py +++ b/magic_pdf/model/sub_modules/table/tablemaster/tablemaster_paddle.py @@ -2,8 +2,8 @@ import os import cv2 import numpy as np -from paddleocr.ppstructure.table.predict_table import TableSystem -from paddleocr.ppstructure.utility import init_args +from ppstructure.table.predict_table import TableSystem +from ppstructure.utility import init_args from PIL import Image from magic_pdf.config.constants import * # noqa: F403 From 023ed9c8b8aa0f2386c1124c82a6444c87a6aae1 Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 01:50:04 +0800 Subject: [PATCH 45/50] build: enable pdfminer.six dependency - Uncomment pdfminer.six in requirements.txt - Specify version 20231228 for pdfminer.six --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index c0e511a5..b7df0569 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,5 +9,5 @@ PyMuPDF>=1.24.9 scikit-learn>=1.0.2 torch>=2.2.2 transformers -# pdfminer.six==20231228 +pdfminer.six==20231228 # The requirements.txt must ensure that only necessary external dependencies are introduced. If there are new dependencies to add, please contact the project administrator. From 9a96362db737d36884b4a9eab951c9a655dfd189 Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 14:31:49 +0800 Subject: [PATCH 46/50] build(deps): update torch and torchvision version requirements - Specify torch==2.3.1 and torchvision==0.18.1 for Windows CUDA installation - Add torch and torchvision version constraints in setup.py: - torch>=2.2.2,<=2.3.1 - torchvision>=0.17.2,<=0.18.1 - Update installation instructions in both English and Chinese README files --- docs/README_Windows_CUDA_Acceleration_en_US.md | 2 +- docs/README_Windows_CUDA_Acceleration_zh_CN.md | 2 +- setup.py | 4 +++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/README_Windows_CUDA_Acceleration_en_US.md b/docs/README_Windows_CUDA_Acceleration_en_US.md index e8252145..7aa96990 100644 --- a/docs/README_Windows_CUDA_Acceleration_en_US.md +++ b/docs/README_Windows_CUDA_Acceleration_en_US.md @@ -65,7 +65,7 @@ If your graphics card has at least 8GB of VRAM, follow these steps to test CUDA- 1. **Overwrite the installation of torch and torchvision** supporting CUDA. ``` - pip install --force-reinstall torch torchvision --index-url https://download.pytorch.org/whl/cu118 + pip install --force-reinstall torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu118 ``` 2. **Modify the value of `"device-mode"`** in the `magic-pdf.json` configuration file located in your user directory. diff --git a/docs/README_Windows_CUDA_Acceleration_zh_CN.md b/docs/README_Windows_CUDA_Acceleration_zh_CN.md index 0b156994..dcdb71ee 100644 --- a/docs/README_Windows_CUDA_Acceleration_zh_CN.md +++ b/docs/README_Windows_CUDA_Acceleration_zh_CN.md @@ -66,7 +66,7 @@ pip install -U magic-pdf[full] --extra-index-url https://wheels.myhloli.com -i h **1.覆盖安装支持cuda的torch和torchvision** ```bash -pip install --force-reinstall torch torchvision --index-url https://download.pytorch.org/whl/cu118 +pip install --force-reinstall torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu118 ``` **2.修改【用户目录】中配置文件magic-pdf.json中"device-mode"的值** diff --git a/setup.py b/setup.py index a1f84a0e..aaa4c1b3 100644 --- a/setup.py +++ b/setup.py @@ -36,7 +36,9 @@ if __name__ == '__main__': "paddlepaddle==3.0.0b1;platform_system=='Linux'", "paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", ], - "full": ["unimernet==0.2.2", # unimernet升级0.2.1 + "full": ["unimernet==0.2.2", # unimernet升级0.2.2,移除torchtext的依赖 + "torch>=2.2.2,<=2.3.1", # torch2.4.0及之后版本未测试，先卡住版本上限 + "torchvision>=0.17.2,<=0.18.1", # torchvision 受torch版本约束 "matplotlib<=3.9.0;platform_system=='Windows'", # 3.9.1及之后不提供windows的预编译包，避免一些没有编译环境的windows设备安装失败 "matplotlib;platform_system=='Linux' or platform_system=='Darwin'", # linux 和 macos 不应限制matplotlib的最高版本，以避免无法更新导致的一些bug "ultralytics>=8.3.48", # yolov8,公式检测 From 4e7511fb866216c48c43755fbcf080ecc1139262 Mon Sep 17 00:00:00 2001 From: icecraft Date: Wed, 11 Dec 2024 16:42:16 +0800 Subject: [PATCH 47/50] fix: dup classify pdf type --- magic_pdf/tools/common.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/magic_pdf/tools/common.py b/magic_pdf/tools/common.py index 4f416f40..d64328c4 100644 --- a/magic_pdf/tools/common.py +++ b/magic_pdf/tools/common.py @@ -123,6 +123,9 @@ def do_parse( formula_enable=formula_enable, table_enable=table_enable, ) + pipe_result = infer_result.pipe_txt_mode( + image_writer, debug_mode=True, lang=lang + ) else: infer_result = ds.apply( doc_analyze, @@ -132,9 +135,9 @@ def do_parse( formula_enable=formula_enable, table_enable=table_enable, ) - pipe_result = infer_result.pipe_auto_mode( - image_writer, debug_mode=True, lang=lang - ) + pipe_result = infer_result.pipe_ocr_mode( + image_writer, debug_mode=True, lang=lang + ) elif parse_method == 'txt': infer_result = ds.apply( From f5d812b3135f9a2268c5e6613fb6f91386a7bdc1 Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 17:28:00 +0800 Subject: [PATCH 48/50] feat(layout): improve layout detection for DocLayout_YOLO model - Implement image cropping and pasting technique to enhance layout detection - Adjust detected polygons to original image coordinates - Add comments for better code readability --- magic_pdf/model/pdf_extract_kit.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 3ab078be..648724d5 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -179,7 +179,25 @@ class CustomPEKModel: layout_res = self.layout_model(image, ignore_catids=[]) elif self.layout_model_name == MODEL_NAME.DocLayout_YOLO: # doclayout_yolo - layout_res = self.layout_model.predict(image) + img_pil = Image.fromarray(image) + width, height = img_pil.size + # logger.info(f'width: {width}, height: {height}') + input_res = {"poly":[0,0,width,0,width,height,0,height]} + new_image, useful_list = crop_img(input_res, img_pil, crop_paste_x=width//2, crop_paste_y=0) + paste_x, paste_y, xmin, ymin, xmax, ymax, new_width, new_height = useful_list + layout_res = self.layout_model.predict(new_image) + for res in layout_res: + p1, p2, p3, p4, p5, p6, p7, p8 = res['poly'] + p1 = p1 - paste_x + xmin + p2 = p2 - paste_y + ymin + p3 = p3 - paste_x + xmin + p4 = p4 - paste_y + ymin + p5 = p5 - paste_x + xmin + p6 = p6 - paste_y + ymin + p7 = p7 - paste_x + xmin + p8 = p8 - paste_y + ymin + res['poly'] = [p1, p2, p3, p4, p5, p6, p7, p8] + layout_cost = round(time.time() - layout_start, 2) logger.info(f'layout detection time: {layout_cost}') From 9efc35ecaac51098e55c0bfcf5ae7c47d7c8b18e Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 18:13:03 +0800 Subject: [PATCH 49/50] refactor(magic_pdf): remove unused import in pdf_parse_union_core_v2.py - Remove unused import of ocr_model_init from magic_pdf.model.sub_modules.model_init - Keep existing functionality and structure intact --- magic_pdf/pdf_parse_union_core_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/magic_pdf/pdf_parse_union_core_v2.py b/magic_pdf/pdf_parse_union_core_v2.py index b226e8d3..8461c4c8 100644 --- a/magic_pdf/pdf_parse_union_core_v2.py +++ b/magic_pdf/pdf_parse_union_core_v2.py @@ -27,7 +27,7 @@ try: except ImportError: pass -from magic_pdf.model.sub_modules.model_init import AtomModelSingleton, ocr_model_init +from magic_pdf.model.sub_modules.model_init import AtomModelSingleton from magic_pdf.para.para_split_v3 import para_split from magic_pdf.pre_proc.construct_page_dict import ocr_construct_page_component_v2 from magic_pdf.pre_proc.cut_image import ocr_cut_image_and_table From ef78819aa9891a6c1feedc5ddbc7b7f3338cad6e Mon Sep 17 00:00:00 2001 From: myhloli Date: Wed, 11 Dec 2024 18:31:51 +0800 Subject: [PATCH 50/50] refactor(draw_bbox): remove redundant '_line_sort' suffix from output filename - Updated the filename generation logic in the draw_bbox function - Removed the unnecessary '_line_sort' suffix from the output PDF filename --- magic_pdf/libs/draw_bbox.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/magic_pdf/libs/draw_bbox.py b/magic_pdf/libs/draw_bbox.py index 7ec359ea..5cf04ac3 100644 --- a/magic_pdf/libs/draw_bbox.py +++ b/magic_pdf/libs/draw_bbox.py @@ -391,7 +391,7 @@ def draw_line_sort_bbox(pdf_info, pdf_bytes, out_path, filename): for i, page in enumerate(pdf_docs): draw_bbox_with_number(i, layout_bbox_list, page, [255, 0, 0], False) - pdf_docs.save(f'{out_path}/{filename}_line_sort.pdf') + pdf_docs.save(f'{out_path}/{filename}') def draw_layout_sort_bbox(pdf_info, pdf_bytes, out_path, filename):