ocr_validator_utils.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966
  1. """
  2. OCR验证工具的工具函数模块
  3. 包含数据处理、图像处理、统计分析等功能
  4. """
  5. import json
  6. import pandas as pd
  7. import numpy as np
  8. from pathlib import Path
  9. from PIL import Image, ImageDraw
  10. from typing import Dict, List, Optional, Tuple, Union
  11. from io import StringIO, BytesIO
  12. import re
  13. from html import unescape
  14. import yaml
  15. import base64
  16. from urllib.parse import urlparse
  17. import cv2
  18. import os
  19. def load_config(config_path: str = "config.yaml") -> Dict:
  20. """加载配置文件"""
  21. try:
  22. with open(config_path, 'r', encoding='utf-8') as f:
  23. return yaml.safe_load(f)
  24. except Exception as e:
  25. # 返回默认配置
  26. return get_default_config()
  27. def get_default_config() -> Dict:
  28. """获取默认配置"""
  29. return {
  30. 'styles': {
  31. 'font_sizes': {'small': 10, 'medium': 12, 'large': 14, 'extra_large': 16},
  32. 'colors': {
  33. 'primary': '#0288d1', 'secondary': '#ff9800', 'success': '#4caf50',
  34. 'error': '#f44336', 'warning': '#ff9800', 'background': '#fafafa', 'text': '#333333'
  35. },
  36. 'layout': {'default_zoom': 1.0, 'default_height': 600, 'sidebar_width': 0.3, 'content_width': 0.7}
  37. },
  38. 'ui': {
  39. 'page_title': 'OCR可视化校验工具', 'page_icon': '🔍', 'layout': 'wide',
  40. 'sidebar_state': 'expanded', 'default_font_size': 'medium', 'default_layout': '标准布局'
  41. },
  42. 'paths': {
  43. 'ocr_out_dir': './sample_data', 'src_img_dir': './sample_data',
  44. 'supported_image_formats': ['.png', '.jpg', '.jpeg']
  45. },
  46. 'ocr': {
  47. 'min_text_length': 2, 'default_confidence': 1.0, 'exclude_texts': ['Picture', ''],
  48. 'tools': {
  49. 'dots_ocr': {
  50. 'name': 'Dots OCR', 'json_structure': 'array',
  51. 'text_field': 'text', 'bbox_field': 'bbox', 'category_field': 'category'
  52. },
  53. 'ppstructv3': {
  54. 'name': 'PPStructV3', 'json_structure': 'object', 'parsing_results_field': 'parsing_res_list',
  55. 'text_field': 'block_content', 'bbox_field': 'block_bbox', 'category_field': 'block_label'
  56. }
  57. },
  58. 'auto_detection': {
  59. 'enabled': True,
  60. 'rules': [
  61. {'field_exists': 'parsing_res_list', 'tool_type': 'ppstructv3'},
  62. {'json_is_array': True, 'tool_type': 'dots_ocr'}
  63. ]
  64. }
  65. }
  66. }
  67. def load_css_styles(css_path: str = "styles.css") -> str:
  68. """加载CSS样式文件"""
  69. try:
  70. with open(css_path, 'r', encoding='utf-8') as f:
  71. return f.read()
  72. except Exception:
  73. # 返回基本样式
  74. return """
  75. .main > div { background-color: white !important; color: #333333 !important; }
  76. .stApp { background-color: white !important; }
  77. .block-container { background-color: white !important; color: #333333 !important; }
  78. """
  79. def rotate_image_and_coordinates(
  80. image: Image.Image,
  81. angle: float,
  82. coordinates_list: List[List[int]],
  83. rotate_coordinates: bool = True
  84. ) -> Tuple[Image.Image, List[List[int]]]:
  85. """
  86. 根据角度旋转图像和坐标 - 修正版本
  87. Args:
  88. image: 原始图像
  89. angle: 旋转角度(度数)
  90. coordinates_list: 坐标列表,每个坐标为[x1, y1, x2, y2]格式
  91. rotate_coordinates: 是否需要旋转坐标(针对不同OCR工具的处理方式)
  92. Returns:
  93. rotated_image: 旋转后的图像
  94. rotated_coordinates: 处理后的坐标列表
  95. """
  96. if angle == 0:
  97. return image, coordinates_list
  98. # 标准化旋转角度
  99. if angle == 270:
  100. rotation_angle = -90 # 顺时针90度
  101. elif angle == 90:
  102. rotation_angle = 90 # 逆时针90度
  103. elif angle == 180:
  104. rotation_angle = 180 # 180度
  105. else:
  106. rotation_angle = angle
  107. # 旋转图像
  108. rotated_image = image.rotate(rotation_angle, expand=True)
  109. # 如果不需要旋转坐标,直接返回原坐标
  110. if not rotate_coordinates:
  111. return rotated_image, coordinates_list
  112. # 获取原始和旋转后的图像尺寸
  113. orig_width, orig_height = image.size
  114. new_width, new_height = rotated_image.size
  115. # 计算旋转后的坐标
  116. rotated_coordinates = []
  117. for coord in coordinates_list:
  118. if len(coord) < 4:
  119. rotated_coordinates.append(coord)
  120. continue
  121. x1, y1, x2, y2 = coord[:4]
  122. # 验证原始坐标是否有效
  123. if x1 < 0 or y1 < 0 or x2 <= x1 or y2 <= y1:
  124. print(f"警告: 无效坐标 {coord}")
  125. rotated_coordinates.append([0, 0, 50, 50]) # 使用默认坐标
  126. continue
  127. # 根据旋转角度变换坐标
  128. if rotation_angle == -90: # 顺时针90度 (270度逆时针)
  129. # 变换公式: (x, y) -> (orig_height - y, x)
  130. new_x1 = orig_height - y2 # 这里是y2
  131. new_y1 = x1
  132. new_x2 = orig_height - y1 # 这里是y1
  133. new_y2 = x2
  134. elif rotation_angle == 90: # 逆时针90度
  135. # 变换公式: (x, y) -> (y, orig_width - x)
  136. new_x1 = y1
  137. new_y1 = orig_width - x2 # 这里是x2
  138. new_x2 = y2
  139. new_y2 = orig_width - x1 # 这里是x1
  140. elif rotation_angle == 180: # 180度
  141. # 变换公式: (x, y) -> (orig_width - x, orig_height - y)
  142. new_x1 = orig_width - x2
  143. new_y1 = orig_height - y2
  144. new_x2 = orig_width - x1
  145. new_y2 = orig_height - y1
  146. else: # 任意角度算法,目前90,-90不对
  147. # 将角度转换为弧度
  148. angle_rad = np.radians(rotation_angle)
  149. cos_angle = np.cos(angle_rad)
  150. sin_angle = np.sin(angle_rad)
  151. # 原图像中心点
  152. orig_center_x = orig_width / 2
  153. orig_center_y = orig_height / 2
  154. # 旋转后图像中心点
  155. new_center_x = new_width / 2
  156. new_center_y = new_height / 2
  157. # 将bbox的四个角点转换为相对于原图像中心的坐标
  158. corners = [
  159. (x1 - orig_center_x, y1 - orig_center_y), # 左上角
  160. (x2 - orig_center_x, y1 - orig_center_y), # 右上角
  161. (x2 - orig_center_x, y2 - orig_center_y), # 右下角
  162. (x1 - orig_center_x, y2 - orig_center_y) # 左下角
  163. ]
  164. # 应用旋转矩阵变换每个角点
  165. rotated_corners = []
  166. for x, y in corners:
  167. # 旋转矩阵: [cos(θ) -sin(θ)] [x]
  168. # [sin(θ) cos(θ)] [y]
  169. rotated_x = x * cos_angle - y * sin_angle
  170. rotated_y = x * sin_angle + y * cos_angle
  171. # 转换回绝对坐标(相对于新图像)
  172. abs_x = rotated_x + new_center_x
  173. abs_y = rotated_y + new_center_y
  174. rotated_corners.append((abs_x, abs_y))
  175. # 从旋转后的四个角点计算新的边界框
  176. x_coords = [corner[0] for corner in rotated_corners]
  177. y_coords = [corner[1] for corner in rotated_corners]
  178. new_x1 = int(min(x_coords))
  179. new_y1 = int(min(y_coords))
  180. new_x2 = int(max(x_coords))
  181. new_y2 = int(max(y_coords))
  182. # 确保坐标在有效范围内
  183. new_x1 = max(0, min(new_width, new_x1))
  184. new_y1 = max(0, min(new_height, new_y1))
  185. new_x2 = max(0, min(new_width, new_x2))
  186. new_y2 = max(0, min(new_height, new_y2))
  187. # 确保x1 < x2, y1 < y2
  188. if new_x1 > new_x2:
  189. new_x1, new_x2 = new_x2, new_x1
  190. if new_y1 > new_y2:
  191. new_y1, new_y2 = new_y2, new_y1
  192. rotated_coordinates.append([new_x1, new_y1, new_x2, new_y2])
  193. return rotated_image, rotated_coordinates
  194. def detect_ocr_tool_type(data: Union[List, Dict], config: Dict) -> str:
  195. """自动检测OCR工具类型"""
  196. if not config['ocr']['auto_detection']['enabled']:
  197. return 'dots_ocr' # 默认类型
  198. rules = config['ocr']['auto_detection']['rules']
  199. for rule in rules:
  200. if 'field_exists' in rule:
  201. field_name = rule['field_exists']
  202. if isinstance(data, dict) and field_name in data:
  203. return rule['tool_type']
  204. if 'json_is_array' in rule:
  205. if rule['json_is_array'] and isinstance(data, list):
  206. return rule['tool_type']
  207. # 默认返回dots_ocr
  208. return 'dots_ocr'
  209. def parse_dots_ocr_data(data: List, config: Dict) -> List[Dict]:
  210. """解析Dots OCR格式的数据"""
  211. tool_config = config['ocr']['tools']['dots_ocr']
  212. parsed_data = []
  213. for item in data:
  214. if not isinstance(item, dict):
  215. continue
  216. # 提取字段
  217. text = item.get(tool_config['text_field'], '')
  218. bbox = item.get(tool_config['bbox_field'], [])
  219. category = item.get(tool_config['category_field'], 'Text')
  220. confidence = item.get(tool_config.get('confidence_field', 'confidence'),
  221. config['ocr']['default_confidence'])
  222. if text and bbox and len(bbox) >= 4:
  223. parsed_data.append({
  224. 'text': str(text).strip(),
  225. 'bbox': bbox[:4], # 确保只取前4个坐标
  226. 'category': category,
  227. 'confidence': confidence,
  228. 'source_tool': 'dots_ocr'
  229. })
  230. return parsed_data
  231. def parse_ppstructv3_data(data: Dict, config: Dict) -> List[Dict]:
  232. """解析PPStructV3格式的数据"""
  233. tool_config = config['ocr']['tools']['ppstructv3']
  234. parsed_data = []
  235. # 获取解析结果列表
  236. parsing_results_field = tool_config['parsing_results_field']
  237. if parsing_results_field not in data:
  238. return parsed_data
  239. parsing_results = data[parsing_results_field]
  240. if not isinstance(parsing_results, list):
  241. return parsed_data
  242. for item in parsing_results:
  243. if not isinstance(item, dict):
  244. continue
  245. # 提取字段
  246. text = item.get(tool_config['text_field'], '')
  247. bbox = item.get(tool_config['bbox_field'], [])
  248. category = item.get(tool_config['category_field'], 'text')
  249. confidence = item.get(tool_config.get('confidence_field', 'confidence'),
  250. config['ocr']['default_confidence'])
  251. if text and bbox and len(bbox) >= 4:
  252. parsed_data.append({
  253. 'text': str(text).strip(),
  254. 'bbox': bbox[:4], # 确保只取前4个坐标
  255. 'category': category,
  256. 'confidence': confidence,
  257. 'source_tool': 'ppstructv3'
  258. })
  259. # 如果有OCR文本识别结果,也添加进来
  260. if 'overall_ocr_res' in data:
  261. ocr_res = data['overall_ocr_res']
  262. if isinstance(ocr_res, dict) and 'rec_texts' in ocr_res and 'rec_boxes' in ocr_res:
  263. texts = ocr_res['rec_texts']
  264. boxes = ocr_res['rec_boxes']
  265. scores = ocr_res.get('rec_scores', [])
  266. for i, (text, box) in enumerate(zip(texts, boxes)):
  267. if text and len(box) >= 4:
  268. confidence = scores[i] if i < len(scores) else config['ocr']['default_confidence']
  269. parsed_data.append({
  270. 'text': str(text).strip(),
  271. 'bbox': box[:4],
  272. 'category': 'OCR_Text',
  273. 'confidence': confidence,
  274. 'source_tool': 'ppstructv3_ocr'
  275. })
  276. return parsed_data
  277. def normalize_ocr_data(raw_data: Union[List, Dict], config: Dict) -> List[Dict]:
  278. """统一不同OCR工具的数据格式"""
  279. # 自动检测OCR工具类型
  280. tool_type = detect_ocr_tool_type(raw_data, config)
  281. if tool_type == 'dots_ocr':
  282. return parse_dots_ocr_data(raw_data, config)
  283. elif tool_type == 'ppstructv3':
  284. return parse_ppstructv3_data(raw_data, config)
  285. else:
  286. raise ValueError(f"不支持的OCR工具类型: {tool_type}")
  287. def get_rotation_angle_from_ppstructv3(data: Dict) -> float:
  288. """从PPStructV3数据中获取旋转角度"""
  289. if 'doc_preprocessor_res' in data:
  290. doc_res = data['doc_preprocessor_res']
  291. if isinstance(doc_res, dict) and 'angle' in doc_res:
  292. return float(doc_res['angle'])
  293. return 0.0
  294. def find_image_in_multiple_locations(img_src: str, json_path: str) -> Optional[str]:
  295. """
  296. 在多个可能的位置查找图片文件
  297. """
  298. json_dir = os.path.dirname(json_path)
  299. # 可能的搜索路径
  300. search_paths = [
  301. # 相对于JSON文件的路径
  302. os.path.join(json_dir, img_src),
  303. # 相对于JSON文件父目录的路径
  304. os.path.join(os.path.dirname(json_dir), img_src),
  305. # imgs目录(常见的图片目录)
  306. os.path.join(json_dir, 'imgs', os.path.basename(img_src)),
  307. os.path.join(os.path.dirname(json_dir), 'imgs', os.path.basename(img_src)),
  308. # images目录
  309. os.path.join(json_dir, 'images', os.path.basename(img_src)),
  310. os.path.join(os.path.dirname(json_dir), 'images', os.path.basename(img_src)),
  311. # 同名目录
  312. os.path.join(json_dir, os.path.splitext(os.path.basename(json_path))[0], os.path.basename(img_src)),
  313. ]
  314. # 如果是绝对路径,也加入搜索
  315. if os.path.isabs(img_src):
  316. search_paths.insert(0, img_src)
  317. # 查找存在的文件
  318. for path in search_paths:
  319. if os.path.exists(path):
  320. return path
  321. return None
  322. def process_html_images(html_content: str, json_path: str) -> str:
  323. """
  324. 处理HTML内容中的图片引用,将本地图片转换为base64 - 增强版
  325. """
  326. import re
  327. # 匹配HTML图片标签: <img src="path" ... />
  328. img_pattern = r'<img\s+[^>]*src\s*=\s*["\']([^"\']+)["\'][^>]*/?>'
  329. def replace_html_image(match):
  330. full_tag = match.group(0)
  331. img_src = match.group(1)
  332. # 如果已经是base64或者网络链接,直接返回
  333. if img_src.startswith('data:image') or img_src.startswith('http'):
  334. return full_tag
  335. # 增强的图片查找
  336. full_img_path = find_image_in_multiple_locations(img_src, json_path)
  337. # 尝试转换为base64
  338. try:
  339. if full_img_path and os.path.exists(full_img_path):
  340. with open(full_img_path, 'rb') as img_file:
  341. img_data = img_file.read()
  342. # 获取文件扩展名确定MIME类型
  343. ext = os.path.splitext(full_img_path)[1].lower()
  344. mime_type = {
  345. '.png': 'image/png',
  346. '.jpg': 'image/jpeg',
  347. '.jpeg': 'image/jpeg',
  348. '.gif': 'image/gif',
  349. '.bmp': 'image/bmp',
  350. '.webp': 'image/webp'
  351. }.get(ext, 'image/jpeg')
  352. # 转换为base64
  353. img_base64 = base64.b64encode(img_data).decode('utf-8')
  354. data_url = f"data:{mime_type};base64,{img_base64}"
  355. # 替换src属性,保持其他属性不变
  356. updated_tag = re.sub(
  357. r'src\s*=\s*["\'][^"\']+["\']',
  358. f'src="{data_url}"',
  359. full_tag
  360. )
  361. return updated_tag
  362. else:
  363. # 文件不存在,显示详细的错误信息
  364. search_info = f"搜索路径: {img_src}"
  365. if full_img_path:
  366. search_info += f" -> {full_img_path}"
  367. error_content = f"""
  368. <div style="
  369. color: #d32f2f;
  370. border: 2px dashed #d32f2f;
  371. padding: 10px;
  372. margin: 10px 0;
  373. border-radius: 5px;
  374. background-color: #ffebee;
  375. text-align: center;
  376. ">
  377. <strong>🖼️ 图片无法加载</strong><br>
  378. <small>原始路径: {img_src}</small><br>
  379. <small>JSON文件: {os.path.basename(json_path)}</small><br>
  380. <em>请检查图片文件是否存在</em>
  381. </div>
  382. """
  383. return error_content
  384. except Exception as e:
  385. # 转换失败,返回错误信息
  386. error_content = f"""
  387. <div style="
  388. color: #f57c00;
  389. border: 2px dashed #f57c00;
  390. padding: 10px;
  391. margin: 10px 0;
  392. border-radius: 5px;
  393. background-color: #fff3e0;
  394. text-align: center;
  395. ">
  396. <strong>⚠️ 图片处理失败</strong><br>
  397. <small>文件: {img_src}</small><br>
  398. <small>错误: {str(e)}</small>
  399. </div>
  400. """
  401. return error_content
  402. # 替换所有HTML图片标签
  403. processed_content = re.sub(img_pattern, replace_html_image, html_content, flags=re.IGNORECASE)
  404. return processed_content
  405. def process_markdown_images(md_content: str, json_path: str) -> str:
  406. """
  407. 处理Markdown中的图片引用,将本地图片转换为base64
  408. """
  409. import re
  410. # 匹配Markdown图片语法: ![alt](path)
  411. img_pattern = r'!\[([^\]]*)\]\(([^)]+)\)'
  412. def replace_image(match):
  413. alt_text = match.group(1)
  414. img_path = match.group(2)
  415. # 如果已经是base64或者网络链接,直接返回
  416. if img_path.startswith('data:image') or img_path.startswith('http'):
  417. return match.group(0)
  418. # 处理相对路径
  419. if not os.path.isabs(img_path):
  420. # 相对于JSON文件的路径
  421. json_dir = os.path.dirname(json_path)
  422. full_img_path = os.path.join(json_dir, img_path)
  423. else:
  424. full_img_path = img_path
  425. # 尝试转换为base64
  426. try:
  427. if os.path.exists(full_img_path):
  428. with open(full_img_path, 'rb') as img_file:
  429. img_data = img_file.read()
  430. # 获取文件扩展名确定MIME类型
  431. ext = os.path.splitext(full_img_path)[1].lower()
  432. mime_type = {
  433. '.png': 'image/png',
  434. '.jpg': 'image/jpeg',
  435. '.jpeg': 'image/jpeg',
  436. '.gif': 'image/gif',
  437. '.bmp': 'image/bmp',
  438. '.webp': 'image/webp'
  439. }.get(ext, 'image/jpeg')
  440. # 转换为base64
  441. img_base64 = base64.b64encode(img_data).decode('utf-8')
  442. data_url = f"data:{mime_type};base64,{img_base64}"
  443. return f'![{alt_text}]({data_url})'
  444. else:
  445. # 文件不存在,返回原始链接但添加警告
  446. return f'![{alt_text} (文件不存在)]({img_path})'
  447. except Exception as e:
  448. # 转换失败,返回原始链接
  449. return f'![{alt_text} (加载失败)]({img_path})'
  450. # 替换所有图片引用
  451. processed_content = re.sub(img_pattern, replace_image, md_content)
  452. return processed_content
  453. def process_all_images_in_content(content: str, json_path: str) -> str:
  454. """
  455. 处理内容中的所有图片引用(包括Markdown和HTML格式)
  456. """
  457. # 先处理HTML图片
  458. content = process_html_images(content, json_path)
  459. # 再处理Markdown图片
  460. content = process_markdown_images(content, json_path)
  461. return content
  462. # 修改 load_ocr_data_file 函数
  463. def load_ocr_data_file(json_path: str, config: Dict) -> Tuple[List, str, str]:
  464. """加载OCR相关数据文件"""
  465. json_file = Path(json_path)
  466. ocr_data = []
  467. md_content = ""
  468. image_path = ""
  469. # 加载JSON数据
  470. try:
  471. with open(json_file, 'r', encoding='utf-8') as f:
  472. raw_data = json.load(f)
  473. # 统一数据格式
  474. ocr_data = normalize_ocr_data(raw_data, config)
  475. # 检查是否需要处理图像旋转
  476. rotation_angle = 0.0
  477. if isinstance(raw_data, dict):
  478. rotation_angle = get_rotation_angle_from_ppstructv3(raw_data)
  479. # 如果有旋转角度,记录下来供后续图像处理使用
  480. if rotation_angle != 0:
  481. for item in ocr_data:
  482. item['rotation_angle'] = rotation_angle
  483. except Exception as e:
  484. raise Exception(f"加载JSON文件失败: {e}")
  485. # 加载MD文件
  486. md_file = json_file.with_suffix('.md')
  487. if md_file.exists():
  488. with open(md_file, 'r', encoding='utf-8') as f:
  489. raw_md_content = f.read()
  490. # 处理内容中的所有图片引用(HTML和Markdown)
  491. md_content = process_all_images_in_content(raw_md_content, str(json_file))
  492. # 推断图片路径
  493. image_name = json_file.stem
  494. src_img_dir = Path(config['paths']['src_img_dir'])
  495. image_candidates = []
  496. for ext in config['paths']['supported_image_formats']:
  497. image_candidates.extend([
  498. src_img_dir / f"{image_name}{ext}",
  499. json_file.parent / f"{image_name}{ext}",
  500. # 对于PPStructV3,可能图片名包含page信息 # 去掉page后缀的通用匹配
  501. src_img_dir / f"{image_name.split('_page_')[0]}{ext}" if '_page_' in image_name else None,
  502. ])
  503. # 移除None值
  504. image_candidates = [candidate for candidate in image_candidates if candidate is not None]
  505. for candidate in image_candidates:
  506. if candidate.exists():
  507. image_path = str(candidate)
  508. break
  509. return ocr_data, md_content, image_path
  510. def process_ocr_data(ocr_data: List, config: Dict) -> Dict[str, List]:
  511. """处理OCR数据,建立文本到bbox的映射"""
  512. text_bbox_mapping = {}
  513. exclude_texts = config['ocr']['exclude_texts']
  514. min_text_length = config['ocr']['min_text_length']
  515. if not isinstance(ocr_data, list):
  516. return text_bbox_mapping
  517. for i, item in enumerate(ocr_data):
  518. if not isinstance(item, dict):
  519. continue
  520. text = str(item['text']).strip()
  521. if text and text not in exclude_texts and len(text) >= min_text_length:
  522. bbox = item['bbox']
  523. if isinstance(bbox, list) and len(bbox) == 4:
  524. if text not in text_bbox_mapping:
  525. text_bbox_mapping[text] = []
  526. text_bbox_mapping[text].append({
  527. 'bbox': bbox,
  528. 'category': item.get('category', 'Text'),
  529. 'index': i,
  530. 'confidence': item.get('confidence', config['ocr']['default_confidence']),
  531. 'source_tool': item.get('source_tool', 'unknown'),
  532. 'rotation_angle': item.get('rotation_angle', 0.0) # 添加旋转角度信息
  533. })
  534. return text_bbox_mapping
  535. def find_available_ocr_files(ocr_out_dir: str) -> List[str]:
  536. """查找可用的OCR文件"""
  537. available_files = []
  538. # 搜索多个可能的目录
  539. search_dirs = [
  540. Path(ocr_out_dir),
  541. ]
  542. for search_dir in search_dirs:
  543. if search_dir.exists():
  544. # 递归搜索JSON文件
  545. for json_file in search_dir.rglob("*.json"):
  546. available_files.append(str(json_file))
  547. # 去重并排序
  548. available_files = sorted(list(set(available_files)))
  549. return available_files
  550. def get_ocr_tool_info(ocr_data: List) -> Dict:
  551. """获取OCR工具信息统计"""
  552. tool_counts = {}
  553. for item in ocr_data:
  554. if isinstance(item, dict):
  555. source_tool = item.get('source_tool', 'unknown')
  556. tool_counts[source_tool] = tool_counts.get(source_tool, 0) + 1
  557. return tool_counts
  558. def draw_bbox_on_image(image: Image.Image, bbox: List[int], color: str = "red", width: int = 3) -> Image.Image:
  559. """在图片上绘制bbox框"""
  560. img_copy = image.copy()
  561. draw = ImageDraw.Draw(img_copy)
  562. x1, y1, x2, y2 = bbox
  563. # 绘制矩形框
  564. draw.rectangle([x1, y1, x2, y2], outline=color, width=width)
  565. # 添加半透明填充
  566. overlay = Image.new('RGBA', img_copy.size, (0, 0, 0, 0))
  567. overlay_draw = ImageDraw.Draw(overlay)
  568. color_map = {
  569. "red": (255, 0, 0, 30),
  570. "blue": (0, 0, 255, 30),
  571. "green": (0, 255, 0, 30)
  572. }
  573. fill_color = color_map.get(color, (255, 255, 0, 30))
  574. overlay_draw.rectangle([x1, y1, x2, y2], fill=fill_color)
  575. img_copy = Image.alpha_composite(img_copy.convert('RGBA'), overlay).convert('RGB')
  576. return img_copy
  577. def get_ocr_statistics(ocr_data: List, text_bbox_mapping: Dict, marked_errors: set) -> Dict:
  578. """获取OCR数据统计信息"""
  579. if not isinstance(ocr_data, list) or not ocr_data:
  580. return {
  581. 'total_texts': 0, 'clickable_texts': 0, 'marked_errors': 0,
  582. 'categories': {}, 'accuracy_rate': 0, 'tool_info': {}
  583. }
  584. total_texts = len(ocr_data)
  585. clickable_texts = len(text_bbox_mapping)
  586. marked_errors_count = len(marked_errors)
  587. # 按类别统计
  588. categories = {}
  589. for item in ocr_data:
  590. if isinstance(item, dict):
  591. category = item.get('category', 'Unknown')
  592. categories[category] = categories.get(category, 0) + 1
  593. # OCR工具信息统计
  594. tool_info = get_ocr_tool_info(ocr_data)
  595. accuracy_rate = (clickable_texts - marked_errors_count) / clickable_texts * 100 if clickable_texts > 0 else 0
  596. return {
  597. 'total_texts': total_texts,
  598. 'clickable_texts': clickable_texts,
  599. 'marked_errors': marked_errors_count,
  600. 'categories': categories,
  601. 'accuracy_rate': accuracy_rate,
  602. 'tool_info': tool_info
  603. }
  604. def convert_html_table_to_markdown(content: str) -> str:
  605. """将HTML表格转换为Markdown表格格式 - 支持横向滚动的增强版本"""
  606. def replace_table(match):
  607. table_html = match.group(0)
  608. # 提取所有行
  609. rows = re.findall(r'<tr[^>]*>(.*?)</tr>', table_html, re.DOTALL | re.IGNORECASE)
  610. if not rows:
  611. return table_html
  612. markdown_rows = []
  613. max_cols = 0
  614. # 处理所有行,找出最大列数
  615. processed_rows = []
  616. for row in rows:
  617. # 提取单元格,支持 th 和 td
  618. cells = re.findall(r'<t[hd][^>]*>(.*?)</t[hd]>', row, re.DOTALL | re.IGNORECASE)
  619. if cells:
  620. clean_cells = []
  621. for cell in cells:
  622. cell_text = re.sub(r'<[^>]+>', '', cell).strip()
  623. cell_text = unescape(cell_text)
  624. # 限制单元格长度,避免表格过宽
  625. if len(cell_text) > 30:
  626. cell_text = cell_text[:27] + "..."
  627. clean_cells.append(cell_text or " ") # 空单元格用空格替代
  628. processed_rows.append(clean_cells)
  629. max_cols = max(max_cols, len(clean_cells))
  630. # 统一所有行的列数
  631. for i, row_cells in enumerate(processed_rows):
  632. while len(row_cells) < max_cols:
  633. row_cells.append(" ")
  634. # 构建Markdown行
  635. markdown_row = '| ' + ' | '.join(row_cells) + ' |'
  636. markdown_rows.append(markdown_row)
  637. # 在第一行后添加分隔符
  638. if i == 0:
  639. separator = '| ' + ' | '.join(['---'] * max_cols) + ' |'
  640. markdown_rows.append(separator)
  641. # 添加滚动提示
  642. if max_cols > 8:
  643. scroll_note = "\n> 📋 **提示**: 此表格列数较多,在某些视图中可能需要横向滚动查看完整内容。\n"
  644. return scroll_note + '\n'.join(markdown_rows) if markdown_rows else table_html
  645. return '\n'.join(markdown_rows) if markdown_rows else table_html
  646. # 替换所有HTML表格
  647. converted = re.sub(r'<table[^>]*>.*?</table>', replace_table, content, flags=re.DOTALL | re.IGNORECASE)
  648. return converted
  649. def parse_html_tables(html_content: str) -> List[pd.DataFrame]:
  650. """解析HTML内容中的表格为DataFrame列表"""
  651. try:
  652. tables = pd.read_html(StringIO(html_content))
  653. return tables if tables else []
  654. except Exception:
  655. return []
  656. def create_dynamic_css(config: Dict, font_size_key: str, height: int) -> str:
  657. """根据配置动态创建CSS样式"""
  658. colors = config['styles']['colors']
  659. font_size = config['styles']['font_sizes'][font_size_key]
  660. return f"""
  661. <style>
  662. .dynamic-content {{
  663. height: {height}px;
  664. font-size: {font_size}px !important;
  665. line-height: 1.4;
  666. background-color: {colors['background']} !important;
  667. color: {colors['text']} !important;
  668. border: 1px solid #ddd;
  669. padding: 10px;
  670. border-radius: 5px;
  671. }}
  672. .highlight-selected {{
  673. background-color: {colors['success']} !important;
  674. color: white !important;
  675. }}
  676. .highlight-error {{
  677. background-color: {colors['error']} !important;
  678. color: white !important;
  679. }}
  680. </style>
  681. """
  682. def export_tables_to_excel(tables: List[pd.DataFrame], filename: str = "ocr_tables.xlsx") -> BytesIO:
  683. """导出表格数据到Excel"""
  684. output = BytesIO()
  685. with pd.ExcelWriter(output, engine='openpyxl') as writer:
  686. for i, table in enumerate(tables):
  687. table.to_excel(writer, sheet_name=f'Table_{i+1}', index=False)
  688. return output
  689. def get_table_statistics(tables: List[pd.DataFrame]) -> List[Dict]:
  690. """获取表格统计信息"""
  691. stats = []
  692. for i, table in enumerate(tables):
  693. numeric_cols = len(table.select_dtypes(include=[np.number]).columns)
  694. stats.append({
  695. 'table_index': i + 1,
  696. 'rows': len(table),
  697. 'columns': len(table.columns),
  698. 'numeric_columns': numeric_cols
  699. })
  700. return stats
  701. def group_texts_by_category(text_bbox_mapping: Dict[str, List]) -> Dict[str, List[str]]:
  702. """按类别对文本进行分组"""
  703. categories = {}
  704. for text, info_list in text_bbox_mapping.items():
  705. category = info_list[0]['category']
  706. if category not in categories:
  707. categories[category] = []
  708. categories[category].append(text)
  709. return categories
  710. def get_ocr_tool_rotation_config(ocr_data: List, config: Dict) -> Dict:
  711. """获取OCR工具的旋转配置"""
  712. if not ocr_data or not isinstance(ocr_data, list):
  713. # 默认配置
  714. return {
  715. 'coordinates_are_pre_rotated': False
  716. }
  717. # 从第一个OCR数据项获取工具类型
  718. first_item = ocr_data[0] if ocr_data else {}
  719. source_tool = first_item.get('source_tool', 'dots_ocr')
  720. # 获取工具配置
  721. tools_config = config.get('ocr', {}).get('tools', {})
  722. if source_tool in tools_config:
  723. tool_config = tools_config[source_tool]
  724. return tool_config.get('rotation', {
  725. 'coordinates_are_pre_rotated': False
  726. })
  727. else:
  728. # 默认配置
  729. return {
  730. 'coordinates_are_pre_rotated': False
  731. }
  732. def detect_image_orientation_by_opencv(image_path: str) -> Dict:
  733. """
  734. 使用OpenCV的文本检测来判断图片方向
  735. """
  736. try:
  737. # 读取图像
  738. image = cv2.imread(image_path)
  739. if image is None:
  740. raise ValueError("无法读取图像文件")
  741. height, width = image.shape[:2]
  742. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  743. # 使用EAST文本检测器或其他方法
  744. # 这里使用简单的边缘检测和轮廓分析
  745. edges = cv2.Canny(gray, 50, 150, apertureSize=3)
  746. # 检测直线
  747. lines = cv2.HoughLines(edges, 1, np.pi/180, threshold=100)
  748. if lines is None:
  749. return {
  750. 'detected_angle': 0.0,
  751. 'confidence': 0.0,
  752. 'method': 'opencv_analysis',
  753. 'message': '未检测到足够的直线特征'
  754. }
  755. # 分析直线角度
  756. angles = []
  757. for rho, theta in lines[:, 0]:
  758. angle = theta * 180 / np.pi
  759. # 将角度标准化到0-180度
  760. if angle > 90:
  761. angle = angle - 180
  762. angles.append(angle)
  763. # 统计主要角度
  764. angle_hist = np.histogram(angles, bins=36, range=(-90, 90))[0]
  765. dominant_angle_idx = np.argmax(angle_hist)
  766. dominant_angle = -90 + dominant_angle_idx * 5 # 每个bin 5度
  767. # 将角度映射到标准旋转角度
  768. if -22.5 <= dominant_angle <= 22.5:
  769. detected_angle = 0.0
  770. elif 22.5 < dominant_angle <= 67.5:
  771. detected_angle = 270.0
  772. elif 67.5 < dominant_angle <= 90 or -90 <= dominant_angle < -67.5:
  773. detected_angle = 90.0
  774. else:
  775. detected_angle = 180.0
  776. confidence = angle_hist[dominant_angle_idx] / len(lines) if len(lines) > 0 else 0.0
  777. return {
  778. 'detected_angle': detected_angle,
  779. 'confidence': min(1.0, confidence),
  780. 'method': 'opencv_analysis',
  781. 'line_count': len(lines),
  782. 'dominant_angle': dominant_angle,
  783. 'message': f'基于{len(lines)}条直线检测到旋转角度: {detected_angle}°'
  784. }
  785. except Exception as e:
  786. return {
  787. 'detected_angle': 0.0,
  788. 'confidence': 0.0,
  789. 'method': 'opencv_analysis',
  790. 'error': str(e),
  791. 'message': f'OpenCV检测过程中发生错误: {str(e)}'
  792. }