From cd11ddcd6b182fc580a2ddf7ded39a9ea36be426 Mon Sep 17 00:00:00 2001 From: xu rui Date: Wed, 18 Dec 2024 15:22:56 +0800 Subject: [PATCH 1/2] docs: make sure the generate process of docs work properly --- magic_pdf/data/dataset.py | 2 +- magic_pdf/operators/__init__.py | 125 +++++++++++++++++++++++++++ magic_pdf/operators/models.py | 4 +- next_docs/en/api/model_operators.rst | 2 +- 4 files changed, 129 insertions(+), 4 deletions(-) diff --git a/magic_pdf/data/dataset.py b/magic_pdf/data/dataset.py index 0ee6e3de..b1eaa259 100644 --- a/magic_pdf/data/dataset.py +++ b/magic_pdf/data/dataset.py @@ -9,7 +9,6 @@ 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 -from magic_pdf.model.sub_modules.language_detection.utils import auto_detect_lang class PageableData(ABC): @@ -149,6 +148,7 @@ class PymuDocDataset(Dataset): if lang == '': self._lang = None elif lang == 'auto': + from magic_pdf.model.sub_modules.language_detection.utils import auto_detect_lang self._lang = auto_detect_lang(bits) logger.info(f"lang: {lang}, detect_lang: {self._lang}") else: diff --git a/magic_pdf/operators/__init__.py b/magic_pdf/operators/__init__.py index e69de29b..70eba663 100644 --- a/magic_pdf/operators/__init__.py +++ b/magic_pdf/operators/__init__.py @@ -0,0 +1,125 @@ +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.operators.pipes 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 + """ + pass + + @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 \ No newline at end of file diff --git a/magic_pdf/operators/models.py b/magic_pdf/operators/models.py index 5c044c71..34cbfe4b 100644 --- a/magic_pdf/operators/models.py +++ b/magic_pdf/operators/models.py @@ -11,9 +11,9 @@ from magic_pdf.libs.draw_bbox import draw_model_bbox from magic_pdf.libs.version import __version__ from magic_pdf.operators.pipes import PipeResult from magic_pdf.pdf_parse_union_core_v2 import pdf_parse_union +from magic_pdf.operators import InferenceResultBase - -class InferenceResult: +class InferenceResult(InferenceResultBase): def __init__(self, inference_results: list, dataset: Dataset): """Initialized method. diff --git a/next_docs/en/api/model_operators.rst b/next_docs/en/api/model_operators.rst index c6900441..7f62632d 100644 --- a/next_docs/en/api/model_operators.rst +++ b/next_docs/en/api/model_operators.rst @@ -2,7 +2,7 @@ Model Api ========== -.. autoclass:: magic_pdf.operators.models.InferenceResult +.. autoclass:: magic_pdf.operators.InferenceResultBase :members: :inherited-members: :show-inheritance: From c968ce860dd284f6814fe69456d6c19d5cdf9a18 Mon Sep 17 00:00:00 2001 From: icecraft Date: Wed, 18 Dec 2024 15:33:41 +0800 Subject: [PATCH 2/2] fix: remove pipe_auto_method --- magic_pdf/operators/__init__.py | 36 ++++----------------------------- 1 file changed, 4 insertions(+), 32 deletions(-) diff --git a/magic_pdf/operators/__init__.py b/magic_pdf/operators/__init__.py index 70eba663..3ad22c35 100644 --- a/magic_pdf/operators/__init__.py +++ b/magic_pdf/operators/__init__.py @@ -1,14 +1,12 @@ -from typing import Callable - from abc import ABC, abstractmethod +from typing import Callable from magic_pdf.data.data_reader_writer import DataWriter from magic_pdf.data.dataset import Dataset from magic_pdf.operators.pipes import PipeResult - __use_inside_model__ = True -__model_mode__ = "full" +__model_mode__ = 'full' class InferenceResultBase(ABC): @@ -22,7 +20,7 @@ class InferenceResultBase(ABC): dataset (Dataset): the dataset related with model inference result """ pass - + @abstractmethod def draw_model(self, file_path: str) -> None: """Draw model inference result. @@ -64,32 +62,6 @@ class InferenceResultBase(ABC): """ 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, @@ -122,4 +94,4 @@ class InferenceResultBase(ABC): debug_mode=False, lang=None, ) -> PipeResult: - pass \ No newline at end of file + pass