__init__.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. from abc import ABC, abstractmethod
  2. from typing import Callable
  3. from magic_pdf.data.data_reader_writer import DataWriter
  4. from magic_pdf.data.dataset import Dataset
  5. from magic_pdf.operators.pipes import PipeResult
  6. __use_inside_model__ = True
  7. __model_mode__ = 'full'
  8. class InferenceResultBase(ABC):
  9. @abstractmethod
  10. def __init__(self, inference_results: list, dataset: Dataset):
  11. """Initialized method.
  12. Args:
  13. inference_results (list): the inference result generated by model
  14. dataset (Dataset): the dataset related with model inference result
  15. """
  16. pass
  17. @abstractmethod
  18. def draw_model(self, file_path: str) -> None:
  19. """Draw model inference result.
  20. Args:
  21. file_path (str): the output file path
  22. """
  23. pass
  24. @abstractmethod
  25. def dump_model(self, writer: DataWriter, file_path: str):
  26. """Dump model inference result to file.
  27. Args:
  28. writer (DataWriter): writer handle
  29. file_path (str): the location of target file
  30. """
  31. pass
  32. @abstractmethod
  33. def get_infer_res(self):
  34. """Get the inference result.
  35. Returns:
  36. list: the inference result generated by model
  37. """
  38. pass
  39. @abstractmethod
  40. def apply(self, proc: Callable, *args, **kwargs):
  41. """Apply callable method which.
  42. Args:
  43. proc (Callable): invoke proc as follows:
  44. proc(inference_result, *args, **kwargs)
  45. Returns:
  46. Any: return the result generated by proc
  47. """
  48. pass
  49. def pipe_txt_mode(
  50. self,
  51. imageWriter: DataWriter,
  52. start_page_id=0,
  53. end_page_id=None,
  54. debug_mode=False,
  55. lang=None,
  56. ) -> PipeResult:
  57. """Post-proc the model inference result, Extract the text using the
  58. third library, such as `pymupdf`
  59. Args:
  60. imageWriter (DataWriter): the image writer handle
  61. start_page_id (int, optional): Defaults to 0. Let user select some pages He/She want to process
  62. end_page_id (int, optional): Defaults to the last page index of dataset. Let user select some pages He/She want to process
  63. debug_mode (bool, optional): Defaults to False. will dump more log if enabled
  64. lang (str, optional): Defaults to None.
  65. Returns:
  66. PipeResult: the result
  67. """
  68. pass
  69. @abstractmethod
  70. def pipe_ocr_mode(
  71. self,
  72. imageWriter: DataWriter,
  73. start_page_id=0,
  74. end_page_id=None,
  75. debug_mode=False,
  76. lang=None,
  77. ) -> PipeResult:
  78. pass