table_template_applier.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974
  1. """
  2. 表格模板应用器
  3. 将人工标注的表格结构应用到其他页面
  4. """
  5. import json
  6. from pathlib import Path
  7. from PIL import Image, ImageDraw
  8. from typing import Dict, List, Tuple, Union, Optional
  9. import numpy as np
  10. import argparse
  11. import sys
  12. # 添加父目录到路径
  13. sys.path.insert(0, str(Path(__file__).parent))
  14. try:
  15. from editor.data_processor import get_structure_from_ocr
  16. from table_line_generator import TableLineGenerator
  17. except ImportError:
  18. from .editor.data_processor import get_structure_from_ocr
  19. from .table_line_generator import TableLineGenerator
  20. class TableTemplateApplier:
  21. """表格模板应用器(混合模式)"""
  22. def __init__(self, template_config_path: str):
  23. """初始化时只提取列信息和表头信息"""
  24. with open(template_config_path, 'r', encoding='utf-8') as f:
  25. self.template = json.load(f)
  26. # ✅ 只提取列宽(固定)
  27. self.col_widths = self.template['col_widths']
  28. # ✅ 计算列的相对位置
  29. self.col_offsets = [0]
  30. for width in self.col_widths:
  31. self.col_offsets.append(self.col_offsets[-1] + width)
  32. # ✅ 提取表头高度(通常固定)
  33. rows = self.template['rows']
  34. if rows:
  35. self.header_height = rows[0]['y_end'] - rows[0]['y_start']
  36. else:
  37. self.header_height = 40
  38. # ✅ 计算数据行高度(用于固定行高模式)
  39. if len(rows) > 1:
  40. data_row_heights = [row['y_end'] - row['y_start'] for row in rows[1:]]
  41. # 使用中位数作为典型行高
  42. self.row_height = int(np.median(data_row_heights)) if data_row_heights else 40
  43. # 兜底行高(同样使用中位数)
  44. self.fallback_row_height = self.row_height
  45. else:
  46. # 如果只有表头,使用默认值
  47. self.row_height = 40
  48. self.fallback_row_height = 40
  49. print(f"\n✅ 加载模板配置:")
  50. print(f" 列数: {len(self.col_widths)}")
  51. print(f" 列宽: {self.col_widths}")
  52. print(f" 表头高度: {self.header_height}px")
  53. print(f" 数据行高: {self.row_height}px (用于固定行高模式)")
  54. print(f" 兜底行高: {self.fallback_row_height}px (OCR失败时使用)")
  55. def detect_table_anchor(self, ocr_data: List[Dict]) -> Tuple[int, int]:
  56. """
  57. 检测表格的锚点位置(表头左上角)
  58. 策略:
  59. 1. 找到Y坐标最小的文本框(表头第一行)
  60. 2. 找到X坐标最小的文本框(第一列)
  61. Args:
  62. ocr_data: OCR识别结果
  63. Returns:
  64. (anchor_x, anchor_y): 表格左上角坐标
  65. """
  66. if not ocr_data:
  67. return (0, 0)
  68. # 找到最小的X和Y坐标
  69. min_x = min(item['bbox'][0] for item in ocr_data)
  70. min_y = min(item['bbox'][1] for item in ocr_data)
  71. return (min_x, min_y)
  72. def detect_table_rows(self, ocr_data: List[Dict], header_y: int) -> int:
  73. """
  74. 检测表格的行数(包括表头)
  75. 策略:
  76. 1. 找到Y坐标最大的文本框
  77. 2. 根据数据行高计算行数
  78. 3. 加上表头行
  79. Args:
  80. ocr_data: OCR识别结果
  81. header_y: 表头起始Y坐标
  82. Returns:
  83. 总行数(包括表头)
  84. """
  85. if not ocr_data:
  86. return 1 # 至少有表头
  87. max_y = max(item['bbox'][3] for item in ocr_data)
  88. # 🔧 计算数据区的高度(排除表头)
  89. data_start_y = header_y + self.header_height
  90. data_height = max_y - data_start_y
  91. # 计算数据行数
  92. num_data_rows = max(int(data_height / self.row_height), 0)
  93. # 总行数 = 1行表头 + n行数据
  94. total_rows = 1 + num_data_rows
  95. print(f"📊 行数计算:")
  96. print(f" 表头Y: {header_y}, 数据区起始Y: {data_start_y}")
  97. print(f" 最大Y: {max_y}, 数据区高度: {data_height}px")
  98. print(f" 数据行数: {num_data_rows}, 总行数: {total_rows}")
  99. return total_rows
  100. def apply_template_fixed(self,
  101. image: Image.Image,
  102. ocr_data: Union[List[Dict], Dict], # 🆕 支持 Dict
  103. anchor_x: int = None,
  104. anchor_y: int = None,
  105. num_rows: int = None,
  106. line_width: int = 2,
  107. line_color: Tuple[int, int, int] = (0, 0, 0)) -> Tuple[Image.Image, Dict]:
  108. """
  109. 将模板应用到图片
  110. Args:
  111. image: 目标图片
  112. ocr_data: OCR识别结果(用于自动检测锚点),可以是列表或完整字典
  113. anchor_x: 表格起始X坐标(None=自动检测)
  114. anchor_y: 表头起始Y坐标(None=自动检测)
  115. num_rows: 总行数(None=自动检测)
  116. line_width: 线条宽度
  117. line_color: 线条颜色
  118. Returns:
  119. 绘制了表格线的图片
  120. """
  121. # 🆕 1. 实例化生成器并进行倾斜校正
  122. ocr_data_dict = {'text_boxes': ocr_data}
  123. # 尝试从 ocr_data 列表中获取角度信息(如果它是从 ocr_data 字典中提取出来的 list)
  124. # 但通常 ocr_data 这里只是 text_boxes 列表。
  125. # 我们需要传递包含 image_rotation_angle 和 skew_angle 的字典。
  126. # 由于调用者可能会传入 list,我们需要检查是否有更多信息。
  127. # 这里假设调用者会在传入 list 前处理好,或者我们在这里无法获取。
  128. # 不过,如果是从 parse_ocr_data 获取的 ocr_data,它应该是个 dict。
  129. # apply_template_fixed 的签名是 ocr_data: List[Dict],这意味着它只接收 text_boxes。
  130. # 这可能是一个问题。我们需要修改调用处或者在这里处理。
  131. # 看看 apply_template_to_single_file 是怎么调用的。
  132. # apply_template_to_single_file:
  133. # text_boxes = ocr_data.get('text_boxes', [])
  134. # applier.apply_template_fixed(image, text_boxes, ...)
  135. # 这样我们就丢失了角度信息。
  136. # 我应该修改 apply_template_fixed 的签名,让它接收 Dict 类型的 ocr_data,或者单独传递角度。
  137. # 为了保持兼容性,我可以修改 apply_template_fixed 内部处理。
  138. # 但最好的方式是让它接收整个 ocr_data 字典,就像 apply_template_hybrid 一样。
  139. # 不过,为了最小化修改,我可以在 apply_template_to_single_file 里把角度传进来?
  140. # 不,那得改很多。
  141. # 让我们看看能不能在 apply_template_fixed 里重新构造 ocr_data_dict。
  142. # 如果传入的 ocr_data 是 list,那我们确实没法知道角度。
  143. # 除非我们改变 apply_template_to_single_file 的调用方式。
  144. # 让我们先修改 apply_template_to_single_file 的调用方式,传整个 ocr_data 进去。
  145. # 但是 apply_template_fixed 的签名明确写了 ocr_data: List[Dict]。
  146. # 既然我正在修改这个文件,我可以改变它的签名。
  147. # 或者,我可以像 apply_template_hybrid 一样,增加一个参数 ocr_data_full: Dict = None
  148. # 实际上,apply_template_hybrid 已经接收 ocr_data_dict: Dict。
  149. # apply_template_fixed 接收 List[Dict]。
  150. # 这是一个不一致的地方。
  151. # 我决定修改 apply_template_fixed 的参数,让它也能利用 TableLineGenerator 进行校正。
  152. # 但是 TableLineGenerator 需要完整的 ocr_data 字典才能读取角度。
  153. # 方案:修改 apply_template_fixed 接收 ocr_data_dict。
  154. # 为了兼容旧代码,如果传入的是 list,就包装一下。
  155. # 但是 Python 类型提示 List[Dict] 和 Dict 是不一样的。
  156. # 我可以把参数名改成 ocr_input,类型 Union[List[Dict], Dict]。
  157. # 或者,既然这是内部使用的工具,我直接修改签名,让它接收 Dict。
  158. # 检查一下是否有其他地方调用这个方法。
  159. # 只在 apply_template_to_single_file 调用了。
  160. # 所以我将修改 apply_template_fixed 接收 ocr_data_dict: Dict。
  161. generator = TableLineGenerator(image, {'text_boxes': ocr_data} if isinstance(ocr_data, list) else ocr_data)
  162. corrected_image, angle = generator.correct_skew()
  163. # 获取角度信息
  164. image_rotation_angle = generator.ocr_data.get('image_rotation_angle', 0.0)
  165. skew_angle = generator.ocr_data.get('skew_angle', 0.0)
  166. if abs(angle) > 0.1 or image_rotation_angle != 0:
  167. print(f"🔄 [TemplateApplier] 自动校正: 旋转={image_rotation_angle}°, 倾斜={skew_angle:.2f}°")
  168. # 更新 OCR 数据(generator 内部已经更新了)
  169. ocr_data = generator.ocr_data.get('text_boxes', [])
  170. # 使用校正后的图片
  171. img_with_lines = corrected_image.copy()
  172. else:
  173. img_with_lines = image.copy()
  174. # 如果是字典,提取 list
  175. if isinstance(ocr_data, dict):
  176. ocr_data = ocr_data.get('text_boxes', [])
  177. draw = ImageDraw.Draw(img_with_lines)
  178. # 🔍 自动检测锚点
  179. if anchor_x is None or anchor_y is None:
  180. detected_x, detected_y = self.detect_table_anchor(ocr_data)
  181. anchor_x = anchor_x or detected_x
  182. anchor_y = anchor_y or detected_y
  183. # 🔍 自动检测行数
  184. if num_rows is None:
  185. num_rows = self.detect_table_rows(ocr_data, anchor_y)
  186. print(f"\n📍 表格锚点: ({anchor_x}, {anchor_y})")
  187. print(f"📊 总行数: {num_rows} (1表头 + {num_rows-1}数据)")
  188. # 🎨 生成横线坐标
  189. horizontal_lines = []
  190. # 第1条线:表头顶部
  191. horizontal_lines.append(anchor_y)
  192. # 第2条线:表头底部/数据区顶部
  193. horizontal_lines.append(anchor_y + self.header_height)
  194. # 后续横线:数据行分隔线
  195. current_y = anchor_y + self.header_height
  196. for i in range(num_rows - 1): # 减1因为表头已经占了1行
  197. current_y += self.row_height
  198. horizontal_lines.append(current_y)
  199. # 🎨 生成竖线坐标
  200. vertical_lines = []
  201. for offset in self.col_offsets:
  202. x = anchor_x + offset
  203. vertical_lines.append(x)
  204. print(f"📏 横线坐标: {horizontal_lines[:3]}... (共{len(horizontal_lines)}条)")
  205. print(f"📏 竖线坐标: {vertical_lines[:3]}... (共{len(vertical_lines)}条)")
  206. # 🖊️ 绘制横线
  207. x_start = vertical_lines[0]
  208. x_end = vertical_lines[-1]
  209. for y in horizontal_lines:
  210. draw.line([(x_start, y), (x_end, y)], fill=line_color, width=line_width)
  211. # 🖊️ 绘制竖线
  212. y_start = horizontal_lines[0]
  213. y_end = horizontal_lines[-1]
  214. for x in vertical_lines:
  215. draw.line([(x, y_start), (x, y_end)], fill=line_color, width=line_width)
  216. print(f"✅ 表格绘制完成: {len(horizontal_lines)}行 × {len(vertical_lines)-1}列")
  217. # 🔑 生成结构信息
  218. structure = self._build_structure(
  219. horizontal_lines,
  220. vertical_lines,
  221. anchor_x,
  222. anchor_y,
  223. mode='fixed',
  224. image_rotation_angle=image_rotation_angle,
  225. skew_angle=skew_angle
  226. )
  227. return img_with_lines, structure
  228. def apply_template_hybrid(self,
  229. image: Image.Image,
  230. ocr_data_dict: Dict,
  231. use_ocr_rows: bool = True,
  232. anchor_x: int = None,
  233. anchor_y: int = None,
  234. y_tolerance: int = 5,
  235. line_width: int = 2,
  236. line_color: Tuple[int, int, int] = (0, 0, 0)) -> Tuple[Image.Image, Dict]:
  237. """
  238. 混合模式:使用模板的列 + OCR的行
  239. Args:
  240. image: 目标图片
  241. ocr_data: OCR识别结果(用于检测行)
  242. use_ocr_rows: 是否使用OCR检测的行(True=自适应行高)
  243. anchor_x: 表格起始X坐标(None=自动检测)
  244. anchor_y: 表头起始Y坐标(None=自动检测)
  245. y_tolerance: Y轴聚类容差(像素)
  246. line_width: 线条宽度
  247. line_color: 线条颜色
  248. Returns:
  249. 绘制了表格线的图片, 结构信息
  250. """
  251. # 🆕 1. 实例化生成器并进行倾斜校正
  252. generator = TableLineGenerator(image, ocr_data_dict)
  253. corrected_image, angle = generator.correct_skew()
  254. # 🆕 获取图片旋转角度
  255. image_rotation_angle = ocr_data_dict.get('image_rotation_angle', 0.0)
  256. skew_angle = ocr_data_dict.get('skew_angle', 0.0)
  257. if abs(angle) > 0.1 or image_rotation_angle != 0:
  258. print(f"🔄 [TemplateApplier] 自动校正: 旋转={image_rotation_angle}°, 倾斜={skew_angle:.2f}°")
  259. # 更新 OCR 数据
  260. ocr_data_dict = generator.ocr_data
  261. # 使用校正后的图片
  262. img_with_lines = corrected_image.copy()
  263. else:
  264. img_with_lines = image.copy()
  265. draw = ImageDraw.Draw(img_with_lines)
  266. ocr_data = ocr_data_dict.get('text_boxes', [])
  267. # 🔍 自动检测锚点
  268. if anchor_x is None or anchor_y is None:
  269. detected_x, detected_y = self.detect_table_anchor(ocr_data)
  270. anchor_x = anchor_x or detected_x
  271. anchor_y = anchor_y or detected_y
  272. print(f"\n📍 表格锚点: ({anchor_x}, {anchor_y})")
  273. # ✅ 竖线:使用模板的列宽(固定)
  274. vertical_lines = [anchor_x + offset for offset in self.col_offsets]
  275. print(f"📏 竖线坐标: {vertical_lines} (使用模板,共{len(vertical_lines)}条)")
  276. # ✅ 横线:根据模式选择
  277. if use_ocr_rows and ocr_data:
  278. horizontal_lines = self._detect_rows_from_ocr(
  279. ocr_data, anchor_y, y_tolerance
  280. )
  281. print(f"📏 横线坐标: 使用OCR检测 (共{len(horizontal_lines)}条,自适应行高)")
  282. else:
  283. num_rows = self.detect_table_rows(ocr_data, anchor_y) if ocr_data else 10
  284. horizontal_lines = self._generate_fixed_rows(anchor_y, num_rows)
  285. print(f"📏 横线坐标: 使用固定行高 (共{len(horizontal_lines)}条)")
  286. # 🖊️ 绘制横线
  287. x_start = vertical_lines[0]
  288. x_end = vertical_lines[-1]
  289. for y in horizontal_lines:
  290. draw.line([(x_start, y), (x_end, y)], fill=line_color, width=line_width)
  291. # 🖊️ 绘制竖线
  292. y_start = horizontal_lines[0]
  293. y_end = horizontal_lines[-1]
  294. for x in vertical_lines:
  295. draw.line([(x, y_start), (x, y_end)], fill=line_color, width=line_width)
  296. print(f"✅ 表格绘制完成: {len(horizontal_lines)}行 × {len(vertical_lines)-1}列")
  297. # 🔑 生成结构信息
  298. structure = self._build_structure(
  299. horizontal_lines,
  300. vertical_lines,
  301. anchor_x,
  302. anchor_y,
  303. mode='hybrid',
  304. image_rotation_angle=image_rotation_angle,
  305. skew_angle=skew_angle
  306. )
  307. return img_with_lines, structure
  308. def _detect_rows_from_ocr(self,
  309. ocr_data: List[Dict],
  310. anchor_y: int,
  311. y_tolerance: int = 5) -> List[int]:
  312. """
  313. 从OCR结果中检测行(自适应行高)
  314. 复用 get_structure_from_ocr 统一接口
  315. Args:
  316. ocr_data: OCR识别结果(MinerU 格式的 text_boxes)
  317. anchor_y: 表格起始Y坐标
  318. y_tolerance: Y轴聚类容差(未使用,保留参数兼容性)
  319. Returns:
  320. 横线 y 坐标列表
  321. """
  322. if not ocr_data:
  323. return [anchor_y, anchor_y + self.header_height]
  324. print(f"\n🔍 OCR行检测 (使用 MinerU 算法):")
  325. print(f" 有效文本框数: {len(ocr_data)}")
  326. # 🔑 验证是否为 MinerU 格式
  327. has_cell_index = any('row' in item and 'col' in item for item in ocr_data)
  328. if not has_cell_index:
  329. print(" ⚠️ 警告: OCR数据不包含 row/col 索引,可能不是 MinerU 格式")
  330. print(" ⚠️ 混合模式需要 MinerU 格式的 JSON 文件")
  331. return [anchor_y, anchor_y + self.header_height]
  332. # 🔑 重构原始数据格式(MinerU 需要完整的 table 结构)
  333. raw_data = {
  334. 'type': 'table',
  335. 'table_cells': ocr_data
  336. }
  337. try:
  338. # ✅ 使用统一接口解析和分析(无需 dummy_image)
  339. table_bbox, structure = get_structure_from_ocr(
  340. raw_data,
  341. tool="mineru"
  342. )
  343. if not structure or 'horizontal_lines' not in structure:
  344. print(" ⚠️ MinerU 分析失败,使用兜底方案")
  345. return [anchor_y, anchor_y + self.header_height]
  346. # 🔑 获取横线坐标
  347. horizontal_lines = structure['horizontal_lines']
  348. # 🔑 调整第一条线到 anchor_y(表头顶部)
  349. if horizontal_lines:
  350. offset = anchor_y - horizontal_lines[0]
  351. horizontal_lines = [y + offset for y in horizontal_lines]
  352. print(f" 检测到行数: {len(horizontal_lines) - 1}")
  353. # 🔑 分析行高分布
  354. if len(horizontal_lines) > 1:
  355. row_heights = []
  356. for i in range(len(horizontal_lines) - 1):
  357. h = horizontal_lines[i+1] - horizontal_lines[i]
  358. row_heights.append(h)
  359. if len(row_heights) > 1:
  360. import numpy as np
  361. print(f" 行高分布: min={min(row_heights)}, "
  362. f"median={int(np.median(row_heights))}, "
  363. f"max={max(row_heights)}")
  364. return horizontal_lines
  365. except Exception as e:
  366. print(f" ⚠️ 解析失败: {e}")
  367. import traceback
  368. traceback.print_exc()
  369. return [anchor_y, anchor_y + self.header_height]
  370. def _generate_fixed_rows(self, anchor_y: int, num_rows: int) -> List[int]:
  371. """生成固定行高的横线(兜底方案)"""
  372. horizontal_lines = [anchor_y]
  373. # 表头
  374. horizontal_lines.append(anchor_y + self.header_height)
  375. # 数据行
  376. current_y = anchor_y + self.header_height
  377. for i in range(num_rows - 1):
  378. current_y += self.fallback_row_height
  379. horizontal_lines.append(current_y)
  380. return horizontal_lines
  381. def _build_structure(self,
  382. horizontal_lines: List[int],
  383. vertical_lines: List[int],
  384. anchor_x: int,
  385. anchor_y: int,
  386. mode: str = 'fixed',
  387. image_rotation_angle: float = 0.0,
  388. skew_angle: float = 0.0) -> Dict:
  389. """构建表格结构信息(统一)"""
  390. # 生成行区间
  391. rows = []
  392. for i in range(len(horizontal_lines) - 1):
  393. rows.append({
  394. 'y_start': horizontal_lines[i],
  395. 'y_end': horizontal_lines[i + 1],
  396. 'bboxes': []
  397. })
  398. # 生成列区间
  399. columns = []
  400. for i in range(len(vertical_lines) - 1):
  401. columns.append({
  402. 'x_start': vertical_lines[i],
  403. 'x_end': vertical_lines[i + 1]
  404. })
  405. # ✅ 根据模式设置正确的 mode 值
  406. if mode == 'hybrid':
  407. mode_value = 'hybrid'
  408. elif mode == 'fixed':
  409. mode_value = 'fixed'
  410. else:
  411. mode_value = mode # 保留原始值
  412. return {
  413. 'rows': rows,
  414. 'columns': columns,
  415. 'horizontal_lines': horizontal_lines,
  416. 'vertical_lines': vertical_lines,
  417. 'col_widths': self.col_widths,
  418. 'row_height': self.row_height if mode == 'fixed' else None,
  419. 'table_bbox': [
  420. vertical_lines[0],
  421. horizontal_lines[0],
  422. vertical_lines[-1],
  423. horizontal_lines[-1]
  424. ],
  425. 'mode': mode_value, # ✅ 确保有 mode 字段
  426. 'anchor': {'x': anchor_x, 'y': anchor_y},
  427. 'modified_h_lines': [], # ✅ 添加修改记录字段
  428. 'modified_v_lines': [], # ✅ 添加修改记录字段
  429. 'image_rotation_angle': image_rotation_angle,
  430. 'skew_angle': skew_angle,
  431. 'is_skew_corrected': abs(skew_angle) > 0.1 or image_rotation_angle != 0
  432. }
  433. def apply_template_to_single_file(
  434. applier: TableTemplateApplier,
  435. image_file: Path,
  436. json_file: Path,
  437. output_dir: Path,
  438. structure_suffix: str = "_structure.json",
  439. use_hybrid_mode: bool = True,
  440. line_width: int = 2,
  441. line_color: Tuple[int, int, int] = (0, 0, 0)
  442. ) -> bool:
  443. """
  444. 应用模板到单个文件
  445. Args:
  446. applier: 模板应用器实例
  447. image_file: 图片文件路径
  448. json_file: OCR JSON文件路径
  449. output_dir: 输出目录
  450. use_hybrid_mode: 是否使用混合模式(需要 MinerU 格式)
  451. line_width: 线条宽度
  452. line_color: 线条颜色
  453. Returns:
  454. 是否成功
  455. """
  456. print(f"📄 处理: {image_file.name}")
  457. try:
  458. # 加载OCR数据
  459. with open(json_file, 'r', encoding='utf-8') as f:
  460. raw_data = json.load(f)
  461. # 🔑 自动检测 OCR 格式
  462. ocr_format = None
  463. if 'parsing_res_list' in raw_data and 'overall_ocr_res' in raw_data:
  464. # PPStructure 格式
  465. ocr_format = 'ppstructure'
  466. elif isinstance(raw_data, (list, dict)):
  467. # 尝试提取 MinerU 格式
  468. table_data = None
  469. if isinstance(raw_data, list):
  470. for item in raw_data:
  471. if isinstance(item, dict) and item.get('type') == 'table':
  472. table_data = item
  473. break
  474. elif isinstance(raw_data, dict) and raw_data.get('type') == 'table':
  475. table_data = raw_data
  476. if table_data and 'table_cells' in table_data:
  477. ocr_format = 'mineru'
  478. else:
  479. raise ValueError("未识别的 OCR 格式")
  480. else:
  481. raise ValueError("未识别的 OCR 格式(仅支持 PPStructure 或 MinerU)")
  482. table_bbox, ocr_data = TableLineGenerator.parse_ocr_data(
  483. raw_data,
  484. tool=ocr_format
  485. )
  486. text_boxes = ocr_data.get('text_boxes', [])
  487. print(f" ✅ 加载OCR数据: {len(text_boxes)} 个文本框")
  488. print(f" 📋 OCR格式: {ocr_format}")
  489. # 加载图片
  490. image = Image.open(image_file)
  491. print(f" ✅ 加载图片: {image.size}")
  492. # 🔑 验证混合模式的格式要求
  493. if use_hybrid_mode and ocr_format != 'mineru':
  494. print(f" ⚠️ 警告: 混合模式需要 MinerU 格式,当前格式为 {ocr_format}")
  495. print(f" ℹ️ 自动切换到完全模板模式")
  496. use_hybrid_mode = False
  497. # 🆕 根据模式选择处理方式
  498. if use_hybrid_mode:
  499. print(f" 🔧 使用混合模式 (模板列 + MinerU 行)")
  500. img_with_lines, structure = applier.apply_template_hybrid(
  501. image,
  502. ocr_data,
  503. use_ocr_rows=True,
  504. line_width=line_width,
  505. line_color=line_color
  506. )
  507. else:
  508. print(f" 🔧 使用完全模板模式 (固定行高)")
  509. img_with_lines, structure = applier.apply_template_fixed(
  510. image,
  511. text_boxes,
  512. line_width=line_width,
  513. line_color=line_color
  514. )
  515. # 保存图片
  516. output_file = output_dir / f"{image_file.stem}.png"
  517. img_with_lines.save(output_file)
  518. # 保存结构配置
  519. structure_file = output_dir / f"{image_file.stem}{structure_suffix}"
  520. with open(structure_file, 'w', encoding='utf-8') as f:
  521. json.dump(structure, f, indent=2, ensure_ascii=False)
  522. print(f" ✅ 保存图片: {output_file.name}")
  523. print(f" ✅ 保存配置: {structure_file.name}")
  524. print(f" 📊 表格: {len(structure['rows'])}行 x {len(structure['columns'])}列")
  525. return True
  526. except Exception as e:
  527. print(f" ❌ 处理失败: {e}")
  528. import traceback
  529. traceback.print_exc()
  530. return False
  531. def apply_template_batch(
  532. template_config_path: str,
  533. image_dir: str,
  534. json_dir: str,
  535. output_dir: str,
  536. structure_suffix: str = "_structure.json",
  537. use_hybrid_mode: bool = False,
  538. line_width: int = 2,
  539. line_color: Tuple[int, int, int] = (0, 0, 0)
  540. ):
  541. """
  542. 批量应用模板到所有图片
  543. Args:
  544. template_config_path: 模板配置路径
  545. image_dir: 图片目录
  546. json_dir: OCR JSON目录
  547. output_dir: 输出目录
  548. line_width: 线条宽度
  549. line_color: 线条颜色
  550. """
  551. applier = TableTemplateApplier(template_config_path)
  552. image_path = Path(image_dir)
  553. json_path = Path(json_dir)
  554. output_path = Path(output_dir)
  555. output_path.mkdir(parents=True, exist_ok=True)
  556. # 查找所有图片
  557. image_files = list(image_path.glob("*.jpg")) + list(image_path.glob("*.png"))
  558. image_files.sort()
  559. print(f"\n🔍 找到 {len(image_files)} 个图片文件")
  560. print(f"📂 图片目录: {image_dir}")
  561. print(f"📂 JSON目录: {json_dir}")
  562. print(f"📂 输出目录: {output_dir}\n")
  563. results = []
  564. success_count = 0
  565. failed_count = 0
  566. for idx, image_file in enumerate(image_files, 1):
  567. print(f"\n{'='*60}")
  568. print(f"[{idx}/{len(image_files)}] 处理: {image_file.name}")
  569. print(f"{'='*60}")
  570. # 查找对应的JSON文件
  571. json_file = json_path / f"{image_file.stem}.json"
  572. if not json_file.exists():
  573. print(f"⚠️ 找不到OCR结果: {json_file.name}")
  574. results.append({
  575. 'source': str(image_file),
  576. 'status': 'skipped',
  577. 'reason': 'no_json'
  578. })
  579. failed_count += 1
  580. continue
  581. if apply_template_to_single_file(
  582. applier, image_file, json_file, output_path, structure_suffix, use_hybrid_mode,
  583. line_width, line_color
  584. ):
  585. results.append({
  586. 'source': str(image_file),
  587. 'json': str(json_file),
  588. 'status': 'success'
  589. })
  590. success_count += 1
  591. else:
  592. results.append({
  593. 'source': str(image_file),
  594. 'json': str(json_file),
  595. 'status': 'error'
  596. })
  597. failed_count += 1
  598. print()
  599. # 保存批处理结果
  600. result_file = output_path / "batch_results.json"
  601. with open(result_file, 'w', encoding='utf-8') as f:
  602. json.dump(results, f, indent=2, ensure_ascii=False)
  603. # 统计
  604. skipped_count = sum(1 for r in results if r['status'] == 'skipped')
  605. print(f"\n{'='*60}")
  606. print(f"🎉 批处理完成!")
  607. print(f"{'='*60}")
  608. print(f"✅ 成功: {success_count}")
  609. print(f"❌ 失败: {failed_count}")
  610. print(f"⚠️ 跳过: {skipped_count}")
  611. print(f"📊 总计: {len(results)}")
  612. print(f"📄 结果保存: {result_file}")
  613. def main():
  614. """主函数"""
  615. parser = argparse.ArgumentParser(
  616. description='应用表格模板到其他页面(支持混合模式)',
  617. formatter_class=argparse.RawDescriptionHelpFormatter,
  618. epilog="""
  619. 示例用法:
  620. 1. 混合模式(推荐,自适应行高):
  621. python table_template_applier.py \\
  622. --template template.json \\
  623. --image-dir /path/to/images \\
  624. --json-dir /path/to/jsons \\
  625. --output-dir /path/to/output \\
  626. --structure-suffix _structure.json \\
  627. --hybrid
  628. 2. 完全模板模式(固定行高):
  629. python table_template_applier.py \\
  630. --template template.json \\
  631. --image-file page.png \\
  632. --json-file page.json \\
  633. --output-dir /path/to/output \\
  634. --structure-suffix _structure.json \\
  635. 模式说明:
  636. - 混合模式(--hybrid): 列宽使用模板,行高根据OCR自适应
  637. - 完全模板模式: 列宽和行高都使用模板(适合固定格式表格)
  638. """
  639. )
  640. # 模板参数
  641. parser.add_argument(
  642. '-t', '--template',
  643. type=str,
  644. required=True,
  645. help='模板配置文件路径(人工标注的第一页结构)'
  646. )
  647. # 文件参数组
  648. file_group = parser.add_argument_group('文件参数(单文件模式)')
  649. file_group.add_argument(
  650. '--image-file',
  651. type=str,
  652. help='图片文件路径'
  653. )
  654. file_group.add_argument(
  655. '--json-file',
  656. type=str,
  657. help='OCR JSON文件路径'
  658. )
  659. # 目录参数组
  660. dir_group = parser.add_argument_group('目录参数(批量模式)')
  661. dir_group.add_argument(
  662. '--image-dir',
  663. type=str,
  664. help='图片目录'
  665. )
  666. dir_group.add_argument(
  667. '--json-dir',
  668. type=str,
  669. help='OCR JSON目录'
  670. )
  671. # 输出参数组
  672. output_group = parser.add_argument_group('输出参数')
  673. output_group.add_argument(
  674. '-o', '--output-dir',
  675. type=str,
  676. required=True,
  677. help='输出目录(必需)'
  678. )
  679. output_group.add_argument(
  680. '--structure-suffix',
  681. type=str,
  682. default='_structure.json',
  683. help='输出结构配置文件后缀(默认: _structure.json)'
  684. )
  685. # 绘图参数组
  686. draw_group = parser.add_argument_group('绘图参数')
  687. draw_group.add_argument(
  688. '-w', '--width',
  689. type=int,
  690. default=2,
  691. help='线条宽度(默认: 2)'
  692. )
  693. draw_group.add_argument(
  694. '-c', '--color',
  695. default='black',
  696. choices=['black', 'blue', 'red'],
  697. help='线条颜色(默认: black)'
  698. )
  699. # 🆕 新增模式参数
  700. mode_group = parser.add_argument_group('模式参数')
  701. mode_group.add_argument(
  702. '--hybrid',
  703. action='store_true',
  704. help='使用混合模式(模板列 + OCR行,自适应行高,推荐)'
  705. )
  706. args = parser.parse_args()
  707. # 颜色映射
  708. color_map = {
  709. 'black': (0, 0, 0),
  710. 'blue': (0, 0, 255),
  711. 'red': (255, 0, 0)
  712. }
  713. line_color = color_map[args.color]
  714. # 验证模板文件
  715. template_path = Path(args.template)
  716. if not template_path.exists():
  717. print(f"❌ 错误: 模板文件不存在: {template_path}")
  718. return
  719. output_path = Path(args.output_dir)
  720. output_path.mkdir(parents=True, exist_ok=True)
  721. # 判断模式
  722. if args.image_file and args.json_file:
  723. # 单文件模式
  724. image_file = Path(args.image_file)
  725. json_file = Path(args.json_file)
  726. if not image_file.exists():
  727. print(f"❌ 错误: 图片文件不存在: {image_file}")
  728. return
  729. if not json_file.exists():
  730. print(f"❌ 错误: JSON文件不存在: {json_file}")
  731. return
  732. print("\n🔧 单文件处理模式")
  733. print(f"📄 模板: {template_path.name}")
  734. print(f"📄 图片: {image_file.name}")
  735. print(f"📄 JSON: {json_file.name}")
  736. print(f"📂 输出: {output_path}\n")
  737. applier = TableTemplateApplier(str(template_path))
  738. success = apply_template_to_single_file(
  739. applier, image_file, json_file, output_path,
  740. use_hybrid_mode=args.hybrid, # 🆕 传递混合模式参数
  741. line_width=args.width,
  742. line_color=line_color
  743. )
  744. if success:
  745. print("\n✅ 处理完成!")
  746. else:
  747. print("\n❌ 处理失败!")
  748. elif args.image_dir and args.json_dir:
  749. # 批量模式
  750. image_dir = Path(args.image_dir)
  751. json_dir = Path(args.json_dir)
  752. if not image_dir.exists():
  753. print(f"❌ 错误: 图片目录不存在: {image_dir}")
  754. return
  755. if not json_dir.exists():
  756. print(f"❌ 错误: JSON目录不存在: {json_dir}")
  757. return
  758. print("\n🔧 批量处理模式")
  759. print(f"📄 模板: {template_path.name}")
  760. apply_template_batch(
  761. str(template_path),
  762. str(image_dir),
  763. str(json_dir),
  764. str(output_path),
  765. structure_suffix=args.structure_suffix,
  766. use_hybrid_mode=args.hybrid, # 🆕 传递混合模式参数
  767. line_width=args.width,
  768. line_color=line_color,
  769. )
  770. else:
  771. parser.print_help()
  772. print("\n❌ 错误: 请指定单文件模式或批量模式的参数")
  773. print("\n提示:")
  774. print(" 单文件模式: --image-file + --json-file")
  775. print(" 批量模式: --image-dir + --json-dir")
  776. if __name__ == "__main__":
  777. print("🚀 启动表格模板批量应用程序...")
  778. import sys
  779. if len(sys.argv) == 1:
  780. # 如果没有命令行参数,使用默认配置运行
  781. print("ℹ️ 未提供命令行参数,使用默认配置运行...")
  782. # 默认配置
  783. default_config = {
  784. "template": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行.wiredtable/康强_北京农村商业银行_page_001_structure.json",
  785. "image-file": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行/康强_北京农村商业银行_page_002.png",
  786. "json-file": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行_page_002.json",
  787. "output-dir": "output/batch_results",
  788. "width": "2",
  789. "color": "black"
  790. }
  791. # default_config = {
  792. # "template": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水.wiredtable/B用户_扫描流水_page_001_structure.json",
  793. # "image-file": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/mineru_vllm_results/B用户_扫描流水/B用户_扫描流水_page_002.png",
  794. # "json-file": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/mineru_vllm_results_cell_bbox/B用户_扫描流水_page_002.json",
  795. # "output-dir": "output/batch_results",
  796. # "width": "2",
  797. # "color": "black"
  798. # }
  799. print("⚙️ 默认参数:")
  800. for key, value in default_config.items():
  801. print(f" --{key}: {value}")
  802. # 构造参数
  803. sys.argv = [sys.argv[0]]
  804. for key, value in default_config.items():
  805. sys.argv.extend([f"--{key}", str(value)])
  806. sys.argv.append("--hybrid") # 使用混合模式
  807. sys.exit(main())