pipeline_analyze.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. import os
  2. import time
  3. import numpy as np
  4. import torch
  5. from pypdfium2 import PdfDocument
  6. from mineru.backend.pipeline.model_init import MineruPipelineModel
  7. from .model_json_to_middle_json import result_to_middle_json
  8. from ...data.data_reader_writer import DataWriter
  9. from ...utils.pdf_classify import classify
  10. from ...utils.pdf_image_tools import load_images_from_pdf
  11. from loguru import logger
  12. from ...utils.model_utils import get_vram, clean_memory
  13. from magic_pdf.libs.config_reader import (get_device, get_formula_config,
  14. get_layout_config,
  15. get_local_models_dir,
  16. get_table_recog_config)
  17. os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1' # 让mps可以fallback
  18. os.environ['NO_ALBUMENTATIONS_UPDATE'] = '1' # 禁止albumentations检查更新
  19. class ModelSingleton:
  20. _instance = None
  21. _models = {}
  22. def __new__(cls, *args, **kwargs):
  23. if cls._instance is None:
  24. cls._instance = super().__new__(cls)
  25. return cls._instance
  26. def get_model(
  27. self,
  28. lang=None,
  29. formula_enable=None,
  30. table_enable=None,
  31. ):
  32. key = (lang, formula_enable, table_enable)
  33. if key not in self._models:
  34. self._models[key] = custom_model_init(
  35. lang=lang,
  36. formula_enable=formula_enable,
  37. table_enable=table_enable,
  38. )
  39. return self._models[key]
  40. def custom_model_init(
  41. lang=None,
  42. formula_enable=None,
  43. table_enable=None,
  44. ):
  45. model_init_start = time.time()
  46. # 从配置文件读取model-dir和device
  47. local_models_dir = get_local_models_dir()
  48. device = get_device()
  49. formula_config = get_formula_config()
  50. if formula_enable is not None:
  51. formula_config['enable'] = formula_enable
  52. table_config = get_table_recog_config()
  53. if table_enable is not None:
  54. table_config['enable'] = table_enable
  55. model_input = {
  56. 'models_dir': local_models_dir,
  57. 'device': device,
  58. 'table_config': table_config,
  59. 'formula_config': formula_config,
  60. 'lang': lang,
  61. }
  62. custom_model = MineruPipelineModel(**model_input)
  63. model_init_cost = time.time() - model_init_start
  64. logger.info(f'model init cost: {model_init_cost}')
  65. return custom_model
  66. def doc_analyze(
  67. pdf_bytes_list,
  68. lang_list,
  69. image_writer: DataWriter | None,
  70. parse_method: str = 'auto',
  71. formula_enable=None,
  72. table_enable=None,
  73. ):
  74. """
  75. 统一处理文档分析函数,根据输入参数类型决定处理单个数据集还是多个数据集
  76. Args:
  77. dataset_or_datasets: 单个Dataset对象或Dataset对象列表
  78. parse_method: 解析方法,'auto'/'ocr'/'txt'
  79. formula_enable: 是否启用公式识别
  80. table_enable: 是否启用表格识别
  81. Returns:
  82. 单个dataset时返回单个model_json,多个dataset时返回model_json列表
  83. """
  84. MIN_BATCH_INFERENCE_SIZE = int(os.environ.get('MINERU_MIN_BATCH_INFERENCE_SIZE', 100))
  85. # 收集所有页面信息
  86. all_pages_info = [] # 存储(dataset_index, page_index, img, ocr, lang, width, height)
  87. all_image_lists = []
  88. all_pdf_docs = []
  89. for pdf_idx, pdf_bytes in enumerate(pdf_bytes_list):
  90. # 确定OCR设置
  91. _ocr = False
  92. if parse_method == 'auto':
  93. if classify(pdf_bytes) == 'ocr':
  94. _ocr = True
  95. elif parse_method == 'ocr':
  96. _ocr = True
  97. _lang = lang_list[pdf_idx]
  98. # 收集每个数据集中的页面
  99. images_list, pdf_doc = load_images_from_pdf(pdf_bytes)
  100. all_image_lists.append(images_list)
  101. all_pdf_docs.append(pdf_doc)
  102. for page_idx in range(len(images_list)):
  103. img_dict = images_list[page_idx]
  104. all_pages_info.append((
  105. pdf_idx, page_idx,
  106. img_dict['img_pil'], _ocr, _lang,
  107. ))
  108. # 准备批处理
  109. images_with_extra_info = [(info[2], info[3], info[4]) for info in all_pages_info]
  110. batch_size = MIN_BATCH_INFERENCE_SIZE
  111. batch_images = [
  112. images_with_extra_info[i:i + batch_size]
  113. for i in range(0, len(images_with_extra_info), batch_size)
  114. ]
  115. # 执行批处理
  116. results = []
  117. processed_images_count = 0
  118. for index, batch_image in enumerate(batch_images):
  119. processed_images_count += len(batch_image)
  120. logger.info(
  121. f'Batch {index + 1}/{len(batch_images)}: '
  122. f'{processed_images_count} pages/{len(images_with_extra_info)} pages'
  123. )
  124. batch_results = may_batch_image_analyze(batch_image, formula_enable, table_enable)
  125. results.extend(batch_results)
  126. # 构建返回结果
  127. # 多数据集模式:按数据集分组结果
  128. infer_results = [[] for _ in datasets]
  129. for i, page_info in enumerate(all_pages_info):
  130. pdf_idx, page_idx, pil_img, _, _ = page_info
  131. result = results[i]
  132. page_info_dict = {'page_no': page_idx, 'width': pil_img.get_width(), 'height': pil_img.get_height()}
  133. page_dict = {'layout_dets': result, 'page_info': page_info_dict}
  134. infer_results[pdf_idx].append(page_dict)
  135. middle_json_list = []
  136. for pdf_idx, model_json in enumerate(infer_results):
  137. images_list = all_image_lists[pdf_idx]
  138. pdf_doc = all_pdf_docs[pdf_idx]
  139. middle_json = result_to_middle_json(model_json, images_list, pdf_doc, image_writer)
  140. middle_json_list.append(middle_json)
  141. return middle_json_list, infer_results
  142. def may_batch_image_analyze(
  143. images_with_extra_info: list[(np.ndarray, bool, str)],
  144. formula_enable=None,
  145. table_enable=None):
  146. # os.environ['CUDA_VISIBLE_DEVICES'] = str(idx)
  147. from .batch_analyze import BatchAnalyze
  148. model_manager = ModelSingleton()
  149. batch_ratio = 1
  150. device = get_device()
  151. if str(device).startswith('npu'):
  152. import torch_npu
  153. if torch_npu.npu.is_available():
  154. torch.npu.set_compile_mode(jit_compile=False)
  155. if str(device).startswith('npu') or str(device).startswith('cuda'):
  156. vram = get_vram(device)
  157. if vram is not None:
  158. gpu_memory = int(os.getenv('VIRTUAL_VRAM_SIZE', round(vram)))
  159. if gpu_memory >= 16:
  160. batch_ratio = 16
  161. elif gpu_memory >= 12:
  162. batch_ratio = 8
  163. elif gpu_memory >= 8:
  164. batch_ratio = 4
  165. elif gpu_memory >= 6:
  166. batch_ratio = 2
  167. else:
  168. batch_ratio = 1
  169. logger.info(f'gpu_memory: {gpu_memory} GB, batch_ratio: {batch_ratio}')
  170. else:
  171. # Default batch_ratio when VRAM can't be determined
  172. batch_ratio = 1
  173. logger.info(f'Could not determine GPU memory, using default batch_ratio: {batch_ratio}')
  174. batch_model = BatchAnalyze(model_manager, batch_ratio, formula_enable, table_enable)
  175. results = batch_model(images_with_extra_info)
  176. clean_memory(get_device())
  177. return results