model_init.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190
  1. import os
  2. import torch
  3. from loguru import logger
  4. from .model_list import AtomicModel
  5. from ...model.layout.doclayout_yolo import DocLayoutYOLOModel
  6. from ...model.mfd.yolo_v8 import YOLOv8MFDModel
  7. from ...model.mfr.unimernet.Unimernet import UnimernetModel
  8. from ...model.ocr.paddleocr2pytorch.pytorch_paddle import PytorchPaddleOCR
  9. from ...model.table.rapid_table import RapidTableModel
  10. doclayout_yolo = "Layout/YOLO/doclayout_yolo_docstructbench_imgsz1280_2501.pt"
  11. yolo_v8_mfd = "MFD/YOLO/yolo_v8_ft.pt"
  12. unimernet_small = "MFR/unimernet_hf_small_2503"
  13. def table_model_init(lang=None):
  14. atom_model_manager = AtomModelSingleton()
  15. ocr_engine = atom_model_manager.get_atom_model(
  16. atom_model_name='ocr',
  17. det_db_box_thresh=0.5,
  18. det_db_unclip_ratio=1.6,
  19. lang=lang
  20. )
  21. table_model = RapidTableModel(ocr_engine)
  22. return table_model
  23. def mfd_model_init(weight, device='cpu'):
  24. if str(device).startswith('npu'):
  25. device = torch.device(device)
  26. mfd_model = YOLOv8MFDModel(weight, device)
  27. return mfd_model
  28. def mfr_model_init(weight_dir, device='cpu'):
  29. mfr_model = UnimernetModel(weight_dir, device)
  30. return mfr_model
  31. def doclayout_yolo_model_init(weight, device='cpu'):
  32. if str(device).startswith('npu'):
  33. device = torch.device(device)
  34. model = DocLayoutYOLOModel(weight, device)
  35. return model
  36. def ocr_model_init(det_db_box_thresh=0.3,
  37. lang=None,
  38. use_dilation=True,
  39. det_db_unclip_ratio=1.8,
  40. ):
  41. if lang is not None and lang != '':
  42. model = PytorchPaddleOCR(
  43. det_db_box_thresh=det_db_box_thresh,
  44. lang=lang,
  45. use_dilation=use_dilation,
  46. det_db_unclip_ratio=det_db_unclip_ratio,
  47. )
  48. else:
  49. model = PytorchPaddleOCR(
  50. det_db_box_thresh=det_db_box_thresh,
  51. use_dilation=use_dilation,
  52. det_db_unclip_ratio=det_db_unclip_ratio,
  53. )
  54. return model
  55. class AtomModelSingleton:
  56. _instance = None
  57. _models = {}
  58. def __new__(cls, *args, **kwargs):
  59. if cls._instance is None:
  60. cls._instance = super().__new__(cls)
  61. return cls._instance
  62. def get_atom_model(self, atom_model_name: str, **kwargs):
  63. lang = kwargs.get('lang', None)
  64. table_model_name = kwargs.get('table_model_name', None)
  65. if atom_model_name in [AtomicModel.OCR]:
  66. key = (atom_model_name, lang)
  67. elif atom_model_name in [AtomicModel.Table]:
  68. key = (atom_model_name, table_model_name, lang)
  69. else:
  70. key = atom_model_name
  71. if key not in self._models:
  72. self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs)
  73. return self._models[key]
  74. def atom_model_init(model_name: str, **kwargs):
  75. atom_model = None
  76. if model_name == AtomicModel.Layout:
  77. atom_model = doclayout_yolo_model_init(
  78. kwargs.get('doclayout_yolo_weights'),
  79. kwargs.get('device')
  80. )
  81. elif model_name == AtomicModel.MFD:
  82. atom_model = mfd_model_init(
  83. kwargs.get('mfd_weights'),
  84. kwargs.get('device')
  85. )
  86. elif model_name == AtomicModel.MFR:
  87. atom_model = mfr_model_init(
  88. kwargs.get('mfr_weight_dir'),
  89. kwargs.get('device')
  90. )
  91. elif model_name == AtomicModel.OCR:
  92. atom_model = ocr_model_init(
  93. kwargs.get('det_db_box_thresh'),
  94. kwargs.get('lang'),
  95. )
  96. elif model_name == AtomicModel.Table:
  97. atom_model = table_model_init(
  98. kwargs.get('lang'),
  99. )
  100. else:
  101. logger.error('model name not allow')
  102. exit(1)
  103. if atom_model is None:
  104. logger.error('model init failed')
  105. exit(1)
  106. else:
  107. return atom_model
  108. class MineruPipelineModel:
  109. def __init__(self, **kwargs):
  110. self.formula_config = kwargs.get('formula_config')
  111. self.apply_formula = self.formula_config.get('enable', True)
  112. self.table_config = kwargs.get('table_config')
  113. self.apply_table = self.table_config.get('enable', True)
  114. self.lang = kwargs.get('lang', None)
  115. self.device = kwargs.get('device', 'cpu')
  116. logger.info(
  117. 'DocAnalysis init, this may take some times......'
  118. )
  119. atom_model_manager = AtomModelSingleton()
  120. models_dir = kwargs.get('models_dir', "")
  121. if not models_dir:
  122. logger.error("can't found models_dir, please set models_dir")
  123. exit(1)
  124. if self.apply_formula:
  125. # 初始化公式检测模型
  126. self.mfd_model = atom_model_manager.get_atom_model(
  127. atom_model_name=AtomicModel.MFD,
  128. mfd_weights=str(
  129. os.path.join(models_dir, yolo_v8_mfd)
  130. ),
  131. device=self.device,
  132. )
  133. # 初始化公式解析模型
  134. mfr_weight_dir = str(
  135. os.path.join(models_dir, unimernet_small)
  136. )
  137. self.mfr_model = atom_model_manager.get_atom_model(
  138. atom_model_name=AtomicModel.MFR,
  139. mfr_weight_dir=mfr_weight_dir,
  140. device=self.device,
  141. )
  142. # 初始化layout模型
  143. self.layout_model = atom_model_manager.get_atom_model(
  144. atom_model_name=AtomicModel.Layout,
  145. doclayout_yolo_weights=str(
  146. os.path.join(models_dir, doclayout_yolo)
  147. ),
  148. device=self.device,
  149. )
  150. # 初始化ocr
  151. self.ocr_model = atom_model_manager.get_atom_model(
  152. atom_model_name=AtomicModel.OCR,
  153. det_db_box_thresh=0.3,
  154. lang=self.lang
  155. )
  156. # init table model
  157. if self.apply_table:
  158. self.table_model = atom_model_manager.get_atom_model(
  159. atom_model_name=AtomicModel.Table,
  160. lang=self.lang,
  161. )
  162. logger.info('DocAnalysis init done!')