batch_analyze.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384
  1. import cv2
  2. from loguru import logger
  3. from tqdm import tqdm
  4. from collections import defaultdict
  5. import numpy as np
  6. from .model_init import AtomModelSingleton
  7. from .model_list import AtomicModel
  8. from ...utils.config_reader import get_formula_enable, get_table_enable
  9. from ...utils.model_utils import crop_img, get_res_list_from_layout_res
  10. from ...utils.ocr_utils import get_adjusted_mfdetrec_res, get_ocr_result_list, OcrConfidence
  11. from ...utils.pdf_image_tools import get_crop_img
  12. YOLO_LAYOUT_BASE_BATCH_SIZE = 1
  13. MFD_BASE_BATCH_SIZE = 1
  14. MFR_BASE_BATCH_SIZE = 16
  15. OCR_DET_BASE_BATCH_SIZE = 16
  16. class BatchAnalyze:
  17. def __init__(self, model_manager, batch_ratio: int, formula_enable, table_enable, enable_ocr_det_batch: bool = True):
  18. self.batch_ratio = batch_ratio
  19. self.formula_enable = get_formula_enable(formula_enable)
  20. self.table_enable = get_table_enable(table_enable)
  21. self.model_manager = model_manager
  22. self.enable_ocr_det_batch = enable_ocr_det_batch
  23. def __call__(self, images_with_extra_info: list) -> list:
  24. if len(images_with_extra_info) == 0:
  25. return []
  26. images_layout_res = []
  27. self.model = self.model_manager.get_model(
  28. lang=None,
  29. formula_enable=self.formula_enable,
  30. table_enable=self.table_enable,
  31. )
  32. atom_model_manager = AtomModelSingleton()
  33. images = [image for image, _, _ in images_with_extra_info]
  34. # doclayout_yolo
  35. layout_images = images.copy()
  36. images_layout_res += self.model.layout_model.batch_predict(
  37. layout_images, YOLO_LAYOUT_BASE_BATCH_SIZE
  38. )
  39. if self.formula_enable:
  40. # 公式检测
  41. images_mfd_res = self.model.mfd_model.batch_predict(
  42. images, MFD_BASE_BATCH_SIZE
  43. )
  44. # 公式识别
  45. images_formula_list = self.model.mfr_model.batch_predict(
  46. images_mfd_res,
  47. images,
  48. batch_size=self.batch_ratio * MFR_BASE_BATCH_SIZE,
  49. )
  50. mfr_count = 0
  51. for image_index in range(len(images)):
  52. images_layout_res[image_index] += images_formula_list[image_index]
  53. mfr_count += len(images_formula_list[image_index])
  54. # 清理显存
  55. # clean_vram(self.model.device, vram_threshold=8)
  56. ocr_res_list_all_page = []
  57. table_res_list_all_page = []
  58. for index in range(len(images)):
  59. _, ocr_enable, _lang = images_with_extra_info[index]
  60. layout_res = images_layout_res[index]
  61. pil_img = images[index]
  62. ocr_res_list, table_res_list, single_page_mfdetrec_res = (
  63. get_res_list_from_layout_res(layout_res)
  64. )
  65. ocr_res_list_all_page.append({'ocr_res_list':ocr_res_list,
  66. 'lang':_lang,
  67. 'ocr_enable':ocr_enable,
  68. 'pil_img':pil_img,
  69. 'single_page_mfdetrec_res':single_page_mfdetrec_res,
  70. 'layout_res':layout_res,
  71. })
  72. for table_res in table_res_list:
  73. # table_img, _ = crop_img(table_res, pil_img)
  74. # bbox = (241, 208, 1475, 2019)
  75. scale = 10/3
  76. crop_xmin, crop_ymin = int(table_res['poly'][0]), int(table_res['poly'][1])
  77. crop_xmax, crop_ymax = int(table_res['poly'][4]), int(table_res['poly'][5])
  78. bbox = (int(crop_xmin/scale), int(crop_ymin/scale), int(crop_xmax/scale), int(crop_ymax/scale))
  79. table_img = get_crop_img(bbox, pil_img, scale=scale)
  80. table_res_list_all_page.append({'table_res':table_res,
  81. 'lang':_lang,
  82. 'table_img':table_img,
  83. })
  84. # OCR检测处理
  85. if self.enable_ocr_det_batch:
  86. # 批处理模式 - 按语言和分辨率分组
  87. # 收集所有需要OCR检测的裁剪图像
  88. all_cropped_images_info = []
  89. for ocr_res_list_dict in ocr_res_list_all_page:
  90. _lang = ocr_res_list_dict['lang']
  91. for res in ocr_res_list_dict['ocr_res_list']:
  92. new_image, useful_list = crop_img(
  93. res, ocr_res_list_dict['pil_img'], crop_paste_x=50, crop_paste_y=50
  94. )
  95. adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(
  96. ocr_res_list_dict['single_page_mfdetrec_res'], useful_list
  97. )
  98. # BGR转换
  99. new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR)
  100. all_cropped_images_info.append((
  101. new_image, useful_list, ocr_res_list_dict, res, adjusted_mfdetrec_res, _lang
  102. ))
  103. # 按语言分组
  104. lang_groups = defaultdict(list)
  105. for crop_info in all_cropped_images_info:
  106. lang = crop_info[5]
  107. lang_groups[lang].append(crop_info)
  108. # 对每种语言按分辨率分组并批处理
  109. for lang, lang_crop_list in lang_groups.items():
  110. if not lang_crop_list:
  111. continue
  112. # logger.info(f"Processing OCR detection for language {lang} with {len(lang_crop_list)} images")
  113. # 获取OCR模型
  114. ocr_model = atom_model_manager.get_atom_model(
  115. atom_model_name=AtomicModel.OCR,
  116. det_db_box_thresh=0.3,
  117. lang=lang
  118. )
  119. # 按分辨率分组并同时完成padding
  120. resolution_groups = defaultdict(list)
  121. for crop_info in lang_crop_list:
  122. cropped_img = crop_info[0]
  123. h, w = cropped_img.shape[:2]
  124. # 使用更大的分组容差,减少分组数量
  125. # 将尺寸标准化到32的倍数
  126. normalized_h = ((h + 32) // 32) * 32 # 向上取整到32的倍数
  127. normalized_w = ((w + 32) // 32) * 32
  128. group_key = (normalized_h, normalized_w)
  129. resolution_groups[group_key].append(crop_info)
  130. # 对每个分辨率组进行批处理
  131. for group_key, group_crops in tqdm(resolution_groups.items(), desc=f"OCR-det {lang}"):
  132. # 计算目标尺寸(组内最大尺寸,向上取整到32的倍数)
  133. max_h = max(crop_info[0].shape[0] for crop_info in group_crops)
  134. max_w = max(crop_info[0].shape[1] for crop_info in group_crops)
  135. target_h = ((max_h + 32 - 1) // 32) * 32
  136. target_w = ((max_w + 32 - 1) // 32) * 32
  137. # 对所有图像进行padding到统一尺寸
  138. batch_images = []
  139. for crop_info in group_crops:
  140. img = crop_info[0]
  141. h, w = img.shape[:2]
  142. # 创建目标尺寸的白色背景
  143. padded_img = np.ones((target_h, target_w, 3), dtype=np.uint8) * 255
  144. # 将原图像粘贴到左上角
  145. padded_img[:h, :w] = img
  146. batch_images.append(padded_img)
  147. # 批处理检测
  148. det_batch_size = min(len(batch_images), self.batch_ratio * OCR_DET_BASE_BATCH_SIZE) # 增加批处理大小
  149. # logger.debug(f"OCR-det batch: {det_batch_size} images, target size: {target_h}x{target_w}")
  150. batch_results = ocr_model.text_detector.batch_predict(batch_images, det_batch_size)
  151. # 处理批处理结果
  152. for i, (crop_info, (dt_boxes, elapse)) in enumerate(zip(group_crops, batch_results)):
  153. new_image, useful_list, ocr_res_list_dict, res, adjusted_mfdetrec_res, _lang = crop_info
  154. if dt_boxes is not None and len(dt_boxes) > 0:
  155. # 直接应用原始OCR流程中的关键处理步骤
  156. from mineru.utils.ocr_utils import (
  157. merge_det_boxes, update_det_boxes, sorted_boxes
  158. )
  159. # 1. 排序检测框
  160. if len(dt_boxes) > 0:
  161. dt_boxes_sorted = sorted_boxes(dt_boxes)
  162. else:
  163. dt_boxes_sorted = []
  164. # 2. 合并相邻检测框
  165. if dt_boxes_sorted:
  166. dt_boxes_merged = merge_det_boxes(dt_boxes_sorted)
  167. else:
  168. dt_boxes_merged = []
  169. # 3. 根据公式位置更新检测框(关键步骤!)
  170. if dt_boxes_merged and adjusted_mfdetrec_res:
  171. dt_boxes_final = update_det_boxes(dt_boxes_merged, adjusted_mfdetrec_res)
  172. else:
  173. dt_boxes_final = dt_boxes_merged
  174. # 构造OCR结果格式
  175. ocr_res = [box.tolist() if hasattr(box, 'tolist') else box for box in dt_boxes_final]
  176. if ocr_res:
  177. ocr_result_list = get_ocr_result_list(
  178. ocr_res, useful_list, ocr_res_list_dict['ocr_enable'], new_image, _lang
  179. )
  180. ocr_res_list_dict['layout_res'].extend(ocr_result_list)
  181. else:
  182. # 原始单张处理模式
  183. for ocr_res_list_dict in tqdm(ocr_res_list_all_page, desc="OCR-det Predict"):
  184. # Process each area that requires OCR processing
  185. _lang = ocr_res_list_dict['lang']
  186. # Get OCR results for this language's images
  187. ocr_model = atom_model_manager.get_atom_model(
  188. atom_model_name=AtomicModel.OCR,
  189. ocr_show_log=False,
  190. det_db_box_thresh=0.3,
  191. lang=_lang
  192. )
  193. for res in ocr_res_list_dict['ocr_res_list']:
  194. new_image, useful_list = crop_img(
  195. res, ocr_res_list_dict['pil_img'], crop_paste_x=50, crop_paste_y=50
  196. )
  197. adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(
  198. ocr_res_list_dict['single_page_mfdetrec_res'], useful_list
  199. )
  200. # OCR-det
  201. new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR)
  202. ocr_res = ocr_model.ocr(
  203. new_image, mfd_res=adjusted_mfdetrec_res, rec=False
  204. )[0]
  205. # Integration results
  206. if ocr_res:
  207. ocr_result_list = get_ocr_result_list(
  208. ocr_res, useful_list, ocr_res_list_dict['ocr_enable'],new_image, _lang
  209. )
  210. ocr_res_list_dict['layout_res'].extend(ocr_result_list)
  211. # 表格识别 table recognition
  212. if self.table_enable:
  213. for table_res_dict in tqdm(table_res_list_all_page, desc="Table Predict"):
  214. _lang = table_res_dict['lang']
  215. # 调整图片方向
  216. img_orientation_cls_model = atom_model_manager.get_atom_model(
  217. atom_model_name=AtomicModel.ImgOrientationCls,
  218. )
  219. try:
  220. rotate_label = img_orientation_cls_model.predict(
  221. table_res_dict["table_img"]
  222. )
  223. except Exception as e:
  224. logger.warning(
  225. f"Image orientation classification failed: {e}, using original image"
  226. )
  227. rotate_label = "0"
  228. np_table_img = np.asarray(table_res_dict["table_img"])
  229. if rotate_label == "270":
  230. np_table_img = cv2.rotate(np_table_img, cv2.ROTATE_90_CLOCKWISE)
  231. elif rotate_label == "90":
  232. np_table_img = cv2.rotate(np_table_img, cv2.ROTATE_90_COUNTERCLOCKWISE)
  233. else:
  234. pass
  235. # 有线表/无线表分类
  236. table_cls_model = atom_model_manager.get_atom_model(
  237. atom_model_name=AtomicModel.TableCls,
  238. )
  239. table_cls_score = 0.5
  240. try:
  241. table_label, table_cls_score = table_cls_model.predict(np_table_img)
  242. except Exception as e:
  243. table_label = AtomicModel.WirelessTable
  244. logger.warning(
  245. f"Table classification failed: {e}, using default model {table_label}"
  246. )
  247. # table_label = AtomicModel.WirelessTable
  248. # logger.debug(f"Table classification result: {table_label}")
  249. if table_label not in [AtomicModel.WirelessTable, AtomicModel.WiredTable]:
  250. raise ValueError(
  251. "Table classification failed, please check the model"
  252. )
  253. # 根据表格分类结果选择有线表格识别模型和无线表格识别模型
  254. table_model = atom_model_manager.get_atom_model(
  255. atom_model_name=table_label,
  256. lang=_lang,
  257. )
  258. html_code, table_cell_bboxes, logic_points, elapse = table_model.predict(np_table_img, table_cls_score)
  259. # 判断是否返回正常
  260. if html_code:
  261. # 检查html_code是否包含'<table>'和'</table>'
  262. if '<table>' in html_code and '</table>' in html_code:
  263. # 选用<table>到</table>的内容,放入table_res_dict['table_res']['html']
  264. start_index = html_code.find('<table>')
  265. end_index = html_code.rfind('</table>') + len('</table>')
  266. table_res_dict['table_res']['html'] = html_code[start_index:end_index]
  267. else:
  268. logger.warning(
  269. 'table recognition processing fails, not found expected HTML table end'
  270. )
  271. else:
  272. logger.warning(
  273. 'table recognition processing fails, not get html return'
  274. )
  275. # Create dictionaries to store items by language
  276. need_ocr_lists_by_lang = {} # Dict of lists for each language
  277. img_crop_lists_by_lang = {} # Dict of lists for each language
  278. for layout_res in images_layout_res:
  279. for layout_res_item in layout_res:
  280. if layout_res_item['category_id'] in [15]:
  281. if 'np_img' in layout_res_item and 'lang' in layout_res_item:
  282. lang = layout_res_item['lang']
  283. # Initialize lists for this language if not exist
  284. if lang not in need_ocr_lists_by_lang:
  285. need_ocr_lists_by_lang[lang] = []
  286. img_crop_lists_by_lang[lang] = []
  287. # Add to the appropriate language-specific lists
  288. need_ocr_lists_by_lang[lang].append(layout_res_item)
  289. img_crop_lists_by_lang[lang].append(layout_res_item['np_img'])
  290. # Remove the fields after adding to lists
  291. layout_res_item.pop('np_img')
  292. layout_res_item.pop('lang')
  293. if len(img_crop_lists_by_lang) > 0:
  294. # Process OCR by language
  295. total_processed = 0
  296. # Process each language separately
  297. for lang, img_crop_list in img_crop_lists_by_lang.items():
  298. if len(img_crop_list) > 0:
  299. # Get OCR results for this language's images
  300. ocr_model = atom_model_manager.get_atom_model(
  301. atom_model_name=AtomicModel.OCR,
  302. det_db_box_thresh=0.3,
  303. lang=lang
  304. )
  305. ocr_res_list = ocr_model.ocr(img_crop_list, det=False, tqdm_enable=True)[0]
  306. # Verify we have matching counts
  307. assert len(ocr_res_list) == len(
  308. need_ocr_lists_by_lang[lang]), f'ocr_res_list: {len(ocr_res_list)}, need_ocr_list: {len(need_ocr_lists_by_lang[lang])} for lang: {lang}'
  309. # Process OCR results for this language
  310. for index, layout_res_item in enumerate(need_ocr_lists_by_lang[lang]):
  311. ocr_text, ocr_score = ocr_res_list[index]
  312. layout_res_item['text'] = ocr_text
  313. layout_res_item['score'] = float(f"{ocr_score:.3f}")
  314. if ocr_score < OcrConfidence.min_confidence:
  315. layout_res_item['category_id'] = 16
  316. else:
  317. layout_res_bbox = [layout_res_item['poly'][0], layout_res_item['poly'][1],
  318. layout_res_item['poly'][4], layout_res_item['poly'][5]]
  319. layout_res_width = layout_res_bbox[2] - layout_res_bbox[0]
  320. layout_res_height = layout_res_bbox[3] - layout_res_bbox[1]
  321. if ocr_text in ['(204号', '(20', '(2', '(2号', '(20号'] and ocr_score < 0.8 and layout_res_width < layout_res_height:
  322. layout_res_item['category_id'] = 16
  323. total_processed += len(img_crop_list)
  324. return images_layout_res