data_processor.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455
  1. """
  2. 数据处理模块
  3. 负责处理 MinerU/PaddleOCR_VL/DotsOCR 数据,添加 bbox 信息
  4. """
  5. from typing import List, Dict, Tuple, Optional
  6. from bs4 import BeautifulSoup
  7. try:
  8. from .text_matcher import TextMatcher
  9. from .bbox_extractor import BBoxExtractor
  10. from .table_cell_matcher import TableCellMatcher
  11. except ImportError:
  12. from text_matcher import TextMatcher
  13. from bbox_extractor import BBoxExtractor
  14. from table_cell_matcher import TableCellMatcher
  15. class DataProcessor:
  16. """数据处理器"""
  17. """_summary_
  18. 1.负责处理 MinerU/PaddleOCR_VL/DotsOCR 数据,添加 table_cells bbox 信息, 其他类型的bbox信息依然使用vl自带的bbox
  19. 2.由于不同OCR工具的输出格式不同,DataProcessor 需要包含多个处理方法,分别处理 MinerU、DotsOCR 和 PaddleOCR_VL 数据, 都先转换成mineru格式再添加table cells bbox信息
  20. 3.使用 TextMatcher 进行文本匹配,TableCellMatcher 进行表单元格匹配
  21. 4.最终输出统一的 MinerU 格式数据
  22. 由于VL模型minerU,dotsocr坐标都是使用的原图坐标,不是旋转后的坐标,PaddleVL使用的时旋转转换后的坐标,而ppstructure使用的ocr文本块是旋转后的坐标,
  23. 因此在处理VL数据时,
  24. 1.首先需要根据ppstructure的旋转角度和原图尺寸,将VL的table坐标转换为旋转后的坐标
  25. 2.通过TableCellMatcher 进行表单元格匹配
  26. 3.再将匹配到的单元格bbox逆向转换为原图坐标,存储在最终输出的MinerU格式数据中
  27. 4.其他类型的bbox信息依然使用vl自带的bbox
  28. """
  29. def __init__(self, text_matcher: TextMatcher, look_ahead_window: int = 10, x_tolerance: int = 3, y_tolerance: int = 10):
  30. """
  31. Args:
  32. text_matcher: 文本匹配器
  33. look_ahead_window: 向前查找窗口
  34. x_tolerance: x轴容差
  35. """
  36. self.text_matcher = text_matcher
  37. self.look_ahead_window = look_ahead_window
  38. # X轴容差, 用于判断文本框是否在同一列
  39. self.x_tolerance = x_tolerance
  40. self.y_tolerance = y_tolerance # Y轴容差, 用于行分组
  41. # 🎯 创建表格单元格匹配器
  42. self.table_cell_matcher = TableCellMatcher(
  43. text_matcher=text_matcher,
  44. x_tolerance=x_tolerance,
  45. y_tolerance=y_tolerance
  46. )
  47. def process_mineru_data(self, mineru_data: List[Dict],
  48. paddle_text_boxes: List[Dict], rotation_angle: float, orig_image_size: Tuple[int, int]) -> List[Dict]:
  49. """
  50. 处理 MinerU 数据,添加 bbox 信息
  51. Args:
  52. mineru_data: MinerU 数据
  53. paddle_text_boxes: PaddleOCR 文字框列表
  54. Returns:
  55. 合并后的数据, table cell使用paddle的bbox,其他类型只是移动指针,bbox还是沿用minerU的bbox
  56. """
  57. merged_data = []
  58. paddle_pointer = 0
  59. last_matched_index = 0
  60. # 按 bbox 排序
  61. mineru_data.sort(
  62. key=lambda x: (x['bbox'][1], x['bbox'][0])
  63. if 'bbox' in x else (float('inf'), float('inf'))
  64. )
  65. for item in mineru_data:
  66. item_type = item.get('type', '')
  67. if item_type == 'table':
  68. if rotation_angle != 0:
  69. inverse_table_bbox = BBoxExtractor.rotate_box_coordinates(item['bbox'], rotation_angle, orig_image_size)
  70. inverse_item = item.copy()
  71. inverse_item['bbox'] = inverse_table_bbox
  72. else:
  73. inverse_item = item
  74. merged_item, paddle_pointer, skew_angle = self._process_table(
  75. inverse_item, paddle_text_boxes, paddle_pointer
  76. )
  77. # 🆕 保存角度信息到表格 item
  78. merged_item['image_rotation_angle'] = rotation_angle # 图片旋转角度
  79. merged_item['skew_angle'] = skew_angle # 倾斜角度
  80. # 如果有旋转,需要将匹配到的单元格bbox逆向转换为原图坐标
  81. if rotation_angle != 0:
  82. for cell in merged_item.get('table_cells', []):
  83. cell_bbox = cell.get('bbox', [])
  84. if cell_bbox:
  85. original_bbox = BBoxExtractor.inverse_rotate_box_coordinates(cell_bbox, rotation_angle, orig_image_size)
  86. cell['bbox'] = original_bbox
  87. merged_item['bbox'] = item['bbox'] # 保持表格的原始bbox不变
  88. merged_data.append(merged_item)
  89. elif item_type in ['text', 'title', 'header', 'footer']:
  90. merged_item, paddle_pointer, last_matched_index = self._process_text(
  91. item, paddle_text_boxes, paddle_pointer, last_matched_index
  92. )
  93. merged_data.append(merged_item)
  94. elif item_type == 'list':
  95. merged_item, paddle_pointer, last_matched_index = self._process_list(
  96. item, paddle_text_boxes, paddle_pointer, last_matched_index
  97. )
  98. merged_data.append(merged_item)
  99. else:
  100. merged_data.append(item.copy())
  101. return merged_data
  102. def process_dotsocr_data(self, dotsocr_data: List[Dict],
  103. paddle_text_boxes: List[Dict],
  104. rotation_angle: float,
  105. orig_image_size: Tuple[int, int]) -> List[Dict]:
  106. """
  107. 处理 DotsOCR 数据(简化版:转换后复用 MinerU 处理逻辑)
  108. Args:
  109. dotsocr_data: DotsOCR 输出数据
  110. paddle_text_boxes: PaddleOCR 文本框
  111. rotation_angle: 旋转角度
  112. orig_image_size: 原始图片尺寸
  113. Returns:
  114. 统一的 MinerU 格式数据(带 table_cells bbox)
  115. """
  116. print(f"📊 处理 DotsOCR 数据: {len(dotsocr_data)} 个块")
  117. # 🎯 第一步:转换为 MinerU 格式
  118. mineru_format_data = []
  119. for item in dotsocr_data:
  120. try:
  121. converted_item = self._convert_dotsocr_to_mineru(item)
  122. if converted_item:
  123. mineru_format_data.append(converted_item)
  124. except Exception as e:
  125. print(f"⚠️ DotsOCR 转换失败: {e}")
  126. continue
  127. print(f" ✓ 转换完成: {len(mineru_format_data)} 个块")
  128. # 🎯 第二步:复用 MinerU 处理逻辑
  129. return self.process_mineru_data(
  130. mineru_data=mineru_format_data,
  131. paddle_text_boxes=paddle_text_boxes,
  132. rotation_angle=rotation_angle,
  133. orig_image_size=orig_image_size
  134. )
  135. def _convert_dotsocr_to_mineru(self, dotsocr_item: Dict) -> Dict:
  136. """
  137. 🎯 将 DotsOCR 格式转换为 MinerU 格式
  138. DotsOCR:
  139. {
  140. "category": "Table",
  141. "bbox": [x1, y1, x2, y2],
  142. "text": "..."
  143. }
  144. MinerU:
  145. {
  146. "type": "table",
  147. "bbox": [x1, y1, x2, y2],
  148. "table_body": "...",
  149. "page_idx": 0
  150. }
  151. """
  152. category = dotsocr_item.get('category', '')
  153. # 🎯 Category 映射
  154. category_map = {
  155. 'Page-header': 'header',
  156. 'Page-footer': 'footer',
  157. 'Picture': 'image',
  158. 'Figure': 'image',
  159. 'Section-header': 'title',
  160. 'Table': 'table',
  161. 'Text': 'text',
  162. 'Title': 'title',
  163. 'List': 'list',
  164. 'Caption': 'title'
  165. }
  166. mineru_type = category_map.get(category, 'text')
  167. # 🎯 基础转换
  168. mineru_item = {
  169. 'type': mineru_type,
  170. 'bbox': dotsocr_item.get('bbox', []),
  171. 'page_idx': 0 # DotsOCR 默认单页
  172. }
  173. # 🎯 处理文本内容
  174. text = dotsocr_item.get('text', '')
  175. if mineru_type == 'table':
  176. # 表格:text -> table_body
  177. mineru_item['table_body'] = text
  178. else:
  179. # 其他类型:保持 text
  180. mineru_item['text'] = text
  181. # 标题级别
  182. if category == 'Section-header':
  183. mineru_item['text_level'] = 1
  184. return mineru_item
  185. def process_paddleocr_vl_data(self, paddleocr_vl_data: Dict,
  186. paddle_text_boxes: List[Dict], rotation_angle: float, orig_image_size: Tuple[int, int]) -> List[Dict]:
  187. """
  188. 处理 PaddleOCR_VL 数据,添加 bbox 信息
  189. Args:
  190. paddleocr_vl_data: PaddleOCR_VL 数据 (JSON 对象)
  191. paddle_text_boxes: PaddleOCR 文字框列表
  192. Returns:
  193. 🎯 MinerU 格式的合并数据(统一输出格式)
  194. """
  195. # 🎯 获取旋转角度和原始图像尺寸
  196. vl_rotation_angle = self._get_rotation_angle_from_vl(paddleocr_vl_data)
  197. vl_orig_image_size = (0,0)
  198. if vl_rotation_angle != 0:
  199. vl_orig_image_size = self._get_original_image_size_from_vl(paddleocr_vl_data)
  200. print(f"🔄 PaddleOCR_VL 检测到旋转角度: {vl_rotation_angle}°")
  201. print(f"📐 原始图像尺寸: {vl_orig_image_size[0]} x {vl_orig_image_size[1]}")
  202. # 提取 parsing_res_list
  203. parsing_res_list = paddleocr_vl_data.get('parsing_res_list', [])
  204. # 按 bbox 排序
  205. parsing_res_list.sort(
  206. key=lambda x: (x['block_bbox'][1], x['block_bbox'][0])
  207. if 'block_bbox' in x else (float('inf'), float('inf'))
  208. )
  209. mineru_format_data = []
  210. for item in parsing_res_list:
  211. # 🎯 先转换 bbox 坐标(如果需要)
  212. if vl_rotation_angle != 0 and orig_image_size:
  213. item = self._transform_vl_block_bbox(item, vl_rotation_angle, vl_orig_image_size)
  214. converted_item = self._convert_paddleocr_vl_to_mineru(item)
  215. if converted_item:
  216. mineru_format_data.append(converted_item)
  217. print(f" ✓ 转换完成: {len(mineru_format_data)} 个块")
  218. # 🎯 第三步:复用 MinerU 处理逻辑
  219. return self.process_mineru_data(
  220. mineru_data=mineru_format_data,
  221. paddle_text_boxes=paddle_text_boxes,
  222. rotation_angle=rotation_angle,
  223. orig_image_size=orig_image_size
  224. )
  225. def _get_rotation_angle_from_vl(self, paddleocr_vl_data: Dict) -> float:
  226. """从 PaddleOCR_VL 数据中获取旋转角度"""
  227. return BBoxExtractor._get_rotation_angle(paddleocr_vl_data)
  228. def _get_original_image_size_from_vl(self, paddleocr_vl_data: Dict) -> tuple:
  229. """从 PaddleOCR_VL 数据中获取原始图像尺寸"""
  230. return BBoxExtractor._get_original_image_size(paddleocr_vl_data)
  231. def _transform_vl_block_bbox(self, item: Dict, angle: float,
  232. orig_image_size: tuple) -> Dict:
  233. """
  234. 转换 PaddleOCR_VL 的 block_bbox 坐标
  235. Args:
  236. item: PaddleOCR_VL 的 block 数据
  237. angle: 旋转角度
  238. orig_image_size: 原始图像尺寸
  239. Returns:
  240. 转换后的 block 数据
  241. """
  242. transformed_item = item.copy()
  243. if 'block_bbox' not in item:
  244. return transformed_item
  245. block_bbox = item['block_bbox']
  246. if len(block_bbox) < 4:
  247. return transformed_item
  248. transformed_bbox = BBoxExtractor.inverse_rotate_box_coordinates(block_bbox, angle, orig_image_size)
  249. transformed_item['block_bbox'] = transformed_bbox
  250. return transformed_item
  251. def _convert_paddleocr_vl_to_mineru(self, paddleocr_vl_item: Dict) -> Dict:
  252. """
  253. 🎯 将 PaddleOCR_VL 格式转换为 MinerU 格式
  254. 基于 PP-DocLayout_plus-L 的 20 种类别
  255. """
  256. block_label = paddleocr_vl_item.get('block_label', '')
  257. # 🎯 PP-DocLayout_plus-L 类别映射(共 20 种)
  258. label_map = {
  259. # 标题类(3种)
  260. 'paragraph_title': 'title',
  261. 'doc_title': 'title',
  262. 'figure_table_chart_title': 'title',
  263. # 文本类(9种)
  264. 'text': 'text',
  265. 'number': 'text',
  266. 'content': 'text',
  267. 'abstract': 'text',
  268. 'footnote': 'text',
  269. 'aside_text': 'text',
  270. 'algorithm': 'text',
  271. 'reference': 'text',
  272. 'reference_content': 'text',
  273. # 页眉页脚(2种)
  274. 'header': 'header',
  275. 'footer': 'footer',
  276. # 表格(1种)
  277. 'table': 'table',
  278. # 图片/图表(3种)
  279. 'image': 'image',
  280. 'chart': 'image',
  281. 'seal': 'image',
  282. # 公式(2种)
  283. 'formula': 'equation',
  284. 'formula_number': 'equation'
  285. }
  286. mineru_type = label_map.get(block_label, 'text')
  287. mineru_item = {
  288. 'type': mineru_type,
  289. 'bbox': paddleocr_vl_item.get('block_bbox', []),
  290. 'page_idx': 0
  291. }
  292. content = paddleocr_vl_item.get('block_content', '')
  293. if mineru_type == 'table':
  294. mineru_item['table_body'] = content
  295. else:
  296. mineru_item['text'] = content
  297. # 标题级别
  298. if block_label == 'doc_title':
  299. mineru_item['text_level'] = 1
  300. elif block_label == 'paragraph_title':
  301. mineru_item['text_level'] = 2
  302. elif block_label == 'figure_table_chart_title':
  303. mineru_item['text_level'] = 3
  304. return mineru_item
  305. def _process_table(self, item: Dict, paddle_text_boxes: List[Dict],
  306. start_pointer: int) -> Tuple[Dict, int, float]:
  307. """
  308. 处理表格类型(MinerU 格式)
  309. 策略:
  310. - 解析 HTML 表格
  311. - 为每个单元格匹配 PaddleOCR 的 bbox
  312. - 返回处理后的表格、新指针位置和倾斜角度
  313. """
  314. skew_angle = 0.0
  315. table_body = item.get('table_body', '')
  316. if not table_body:
  317. print(f"⚠️ 表格内容为空,跳过")
  318. return item, start_pointer, skew_angle
  319. try:
  320. # 🔑 传入 table_bbox 用于筛选
  321. table_bbox = item.get('bbox') # MinerU 提供的表格边界
  322. # 🎯 委托给 TableCellMatcher
  323. enhanced_html, cells, new_pointer, skew_angle = \
  324. self.table_cell_matcher.enhance_table_html_with_bbox(
  325. table_body,
  326. paddle_text_boxes,
  327. start_pointer,
  328. table_bbox
  329. )
  330. # 更新 item
  331. item['table_body'] = enhanced_html
  332. item['table_cells'] = cells
  333. # 统计信息
  334. matched_count = len(cells)
  335. total_cells = len(BeautifulSoup(table_body, 'html.parser').find_all(['td', 'th']))
  336. print(f" 表格单元格: {matched_count}/{total_cells} 匹配")
  337. return item, new_pointer, skew_angle # 🆕 返回倾斜角度
  338. except Exception as e:
  339. print(f"⚠️ 表格处理失败: {e}")
  340. import traceback
  341. traceback.print_exc()
  342. return item, start_pointer, skew_angle
  343. def _process_text(self, item: Dict, paddle_text_boxes: List[Dict],
  344. paddle_pointer: int, last_matched_index: int) -> Tuple[Dict, int, int]:
  345. """处理文本"""
  346. merged_item = item.copy()
  347. text = item.get('text', '')
  348. matched_bbox, paddle_pointer, last_matched_index = \
  349. self.text_matcher.find_matching_bbox(
  350. text, paddle_text_boxes, paddle_pointer, last_matched_index,
  351. self.look_ahead_window
  352. )
  353. if matched_bbox:
  354. matched_bbox['used'] = True
  355. return merged_item, paddle_pointer, last_matched_index
  356. def _process_list(self, item: Dict, paddle_text_boxes: List[Dict],
  357. paddle_pointer: int, last_matched_index: int) -> Tuple[Dict, int, int]:
  358. """处理列表"""
  359. merged_item = item.copy()
  360. list_items = item.get('list_items', [])
  361. for list_item in list_items:
  362. matched_bbox, paddle_pointer, last_matched_index = \
  363. self.text_matcher.find_matching_bbox(
  364. list_item, paddle_text_boxes, paddle_pointer, last_matched_index,
  365. self.look_ahead_window
  366. )
  367. if matched_bbox:
  368. matched_bbox['used'] = True
  369. return merged_item, paddle_pointer, last_matched_index