visualization_utils.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436
  1. """
  2. 可视化工具模块
  3. 提供文档处理结果的可视化功能:
  4. - Layout 布局可视化
  5. - OCR 结果可视化
  6. - 图片元素保存
  7. """
  8. from pathlib import Path
  9. from typing import Dict, Any, List, Tuple
  10. import numpy as np
  11. from PIL import Image, ImageDraw, ImageFont
  12. import cv2
  13. from loguru import logger
  14. class VisualizationUtils:
  15. """可视化工具类"""
  16. # 颜色映射(与 MinerU BlockType / EnhancedDocPipeline 类别保持一致)
  17. COLOR_MAP = {
  18. # 文本类元素 (TEXT_CATEGORIES)
  19. 'title': (102, 102, 255), # 蓝色
  20. 'text': (153, 0, 76), # 深红
  21. 'ocr_text': (153, 0, 76), # 深红(同 text)
  22. 'low_score_text': (200, 100, 100), # 浅红
  23. 'header': (128, 128, 128), # 灰色
  24. 'footer': (128, 128, 128), # 灰色
  25. 'page_number': (160, 160, 160), # 浅灰
  26. 'ref_text': (180, 180, 180), # 浅灰
  27. 'aside_text': (180, 180, 180), # 浅灰
  28. 'page_footnote': (200, 200, 200), # 浅灰
  29. # 表格相关元素
  30. 'table': (204, 204, 0), # 黄色
  31. 'table_body': (204, 204, 0), # 黄色
  32. 'table_caption': (255, 255, 102), # 浅黄
  33. 'table_footnote': (229, 255, 204), # 浅黄绿
  34. # 图片相关元素
  35. 'image': (153, 255, 51), # 绿色
  36. 'image_body': (153, 255, 51), # 绿色
  37. 'figure': (153, 255, 51), # 绿色
  38. 'image_caption': (102, 178, 255), # 浅蓝
  39. 'image_footnote': (255, 178, 102), # 橙色
  40. # 公式类元素
  41. 'interline_equation': (0, 255, 0), # 亮绿
  42. 'inline_equation': (0, 200, 0), # 绿色
  43. 'equation': (0, 220, 0), # 绿色
  44. 'interline_equation_yolo': (0, 180, 0),
  45. 'interline_equation_number': (0, 160, 0),
  46. # 代码类元素
  47. 'code': (102, 0, 204), # 紫色
  48. 'code_body': (102, 0, 204), # 紫色
  49. 'code_caption': (153, 51, 255), # 浅紫
  50. 'algorithm': (128, 0, 255), # 紫色
  51. # 列表类元素
  52. 'list': (40, 169, 92), # 青绿
  53. 'index': (60, 180, 100), # 青绿
  54. # 丢弃类元素
  55. 'abandon': (100, 100, 100), # 深灰
  56. 'discarded': (100, 100, 100), # 深灰
  57. # 错误
  58. 'error': (255, 0, 0), # 红色
  59. }
  60. # OCR 框颜色
  61. OCR_BOX_COLOR = (0, 255, 0) # 绿色
  62. CELL_BOX_COLOR = (255, 165, 0) # 橙色
  63. DISCARD_COLOR = (128, 128, 128) # 灰色
  64. @staticmethod
  65. def save_image_elements(
  66. results: Dict[str, Any],
  67. images_dir: Path,
  68. doc_name: str,
  69. is_pdf: bool = True
  70. ) -> List[str]:
  71. """
  72. 保存图片元素
  73. 命名规则:
  74. - PDF输入: 文件名_page_001_image_1.png
  75. - 图片输入(单页): 文件名_image_1.png
  76. Args:
  77. results: 处理结果
  78. images_dir: 图片输出目录
  79. doc_name: 文档名称
  80. is_pdf: 是否为 PDF 输入
  81. Returns:
  82. 保存的图片路径列表
  83. """
  84. saved_paths = []
  85. image_count = 0
  86. total_pages = len(results.get('pages', []))
  87. for page in results.get('pages', []):
  88. page_idx = page.get('page_idx', 0)
  89. for element in page.get('elements', []):
  90. if element.get('type') in ['image', 'image_body', 'figure']:
  91. content = element.get('content', {})
  92. image_data = content.get('image_data')
  93. if image_data is not None:
  94. image_count += 1
  95. # 根据输入类型决定命名
  96. if is_pdf or total_pages > 1:
  97. image_filename = f"{doc_name}_page_{page_idx + 1}_image_{image_count}.png"
  98. else:
  99. image_filename = f"{doc_name}_image_{image_count}.png"
  100. image_path = images_dir / image_filename
  101. try:
  102. if isinstance(image_data, np.ndarray):
  103. cv2.imwrite(str(image_path), image_data)
  104. else:
  105. Image.fromarray(image_data).save(image_path)
  106. # 更新路径(只保存文件名)
  107. content['image_path'] = image_filename
  108. content.pop('image_data', None)
  109. saved_paths.append(str(image_path))
  110. logger.debug(f"🖼️ Image saved: {image_path}")
  111. except Exception as e:
  112. logger.warning(f"Failed to save image: {e}")
  113. if image_count > 0:
  114. logger.info(f"🖼️ {image_count} images saved to: {images_dir}")
  115. return saved_paths
  116. @staticmethod
  117. def save_layout_images(
  118. results: Dict[str, Any],
  119. output_dir: Path,
  120. doc_name: str,
  121. draw_type_label: bool = True,
  122. draw_bbox_number: bool = True,
  123. is_pdf: bool = True
  124. ) -> List[str]:
  125. """
  126. 保存 Layout 可视化图片
  127. 命名规则:
  128. - PDF输入: 文件名_page_001_layout.png
  129. - 图片输入(单页): 文件名_layout.png
  130. Args:
  131. results: 处理结果
  132. output_dir: 输出目录
  133. doc_name: 文档名称
  134. draw_type_label: 是否绘制类型标签
  135. draw_bbox_number: 是否绘制序号
  136. is_pdf: 是否为 PDF 输入
  137. Returns:
  138. 保存的图片路径列表
  139. """
  140. layout_paths = []
  141. total_pages = len(results.get('pages', []))
  142. for page in results.get('pages', []):
  143. page_idx = page.get('page_idx', 0)
  144. processed_image = page.get('original_image')
  145. if processed_image is None:
  146. processed_image = page.get('processed_image')
  147. if processed_image is None:
  148. logger.warning(f"Page {page_idx}: No image data found for layout visualization")
  149. continue
  150. if isinstance(processed_image, np.ndarray):
  151. image = Image.fromarray(processed_image).convert('RGB')
  152. elif isinstance(processed_image, Image.Image):
  153. image = processed_image.convert('RGB')
  154. else:
  155. continue
  156. draw = ImageDraw.Draw(image, 'RGBA')
  157. font = VisualizationUtils._get_font(14)
  158. # 绘制普通元素
  159. for idx, element in enumerate(page.get('elements', []), 1):
  160. elem_type = element.get('type', '')
  161. bbox = element.get('bbox', [0, 0, 0, 0])
  162. if len(bbox) < 4:
  163. continue
  164. x0, y0, x1, y1 = map(int, bbox[:4])
  165. color = VisualizationUtils.COLOR_MAP.get(elem_type, (255, 0, 0))
  166. # 半透明填充
  167. overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
  168. overlay_draw = ImageDraw.Draw(overlay)
  169. overlay_draw.rectangle([x0, y0, x1, y1], fill=(*color, 50))
  170. image = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
  171. draw = ImageDraw.Draw(image)
  172. # 边框
  173. draw.rectangle([x0, y0, x1, y1], outline=color, width=2)
  174. # 类型标签
  175. if draw_type_label:
  176. label = elem_type.replace('_', ' ').title()
  177. bbox_label = draw.textbbox((x0 + 2, y0 + 2), label, font=font)
  178. draw.rectangle(bbox_label, fill=color)
  179. draw.text((x0 + 2, y0 + 2), label, fill='white', font=font)
  180. # 序号
  181. if draw_bbox_number:
  182. number_text = str(idx)
  183. bbox_number = draw.textbbox((x1 - 25, y0 + 2), number_text, font=font)
  184. draw.rectangle(bbox_number, fill=(255, 0, 0))
  185. draw.text((x1 - 25, y0 + 2), number_text, fill='white', font=font)
  186. # 绘制丢弃元素(灰色样式)
  187. for idx, element in enumerate(page.get('discarded_blocks', []), 1):
  188. original_category = element.get('original_category', 'unknown')
  189. bbox = element.get('bbox', [0, 0, 0, 0])
  190. if len(bbox) < 4:
  191. continue
  192. x0, y0, x1, y1 = map(int, bbox[:4])
  193. # 半透明填充
  194. overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
  195. overlay_draw = ImageDraw.Draw(overlay)
  196. overlay_draw.rectangle([x0, y0, x1, y1], fill=(*VisualizationUtils.DISCARD_COLOR, 30))
  197. image = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
  198. draw = ImageDraw.Draw(image)
  199. # 灰色边框
  200. draw.rectangle([x0, y0, x1, y1], outline=VisualizationUtils.DISCARD_COLOR, width=1)
  201. # 类型标签
  202. if draw_type_label:
  203. label = f"D:{original_category}"
  204. bbox_label = draw.textbbox((x0 + 2, y0 + 2), label, font=font)
  205. draw.rectangle(bbox_label, fill=VisualizationUtils.DISCARD_COLOR)
  206. draw.text((x0 + 2, y0 + 2), label, fill='white', font=font)
  207. # 根据输入类型决定命名
  208. if is_pdf or total_pages > 1:
  209. layout_path = output_dir / f"{doc_name}_page_{page_idx + 1:03d}_layout.png"
  210. else:
  211. layout_path = output_dir / f"{doc_name}_layout.png"
  212. image.save(layout_path)
  213. layout_paths.append(str(layout_path))
  214. logger.info(f"🖼️ Layout image saved: {layout_path}")
  215. return layout_paths
  216. @staticmethod
  217. def save_ocr_images(
  218. results: Dict[str, Any],
  219. output_dir: Path,
  220. doc_name: str,
  221. is_pdf: bool = True
  222. ) -> List[str]:
  223. """
  224. 保存 OCR 可视化图片
  225. 命名规则:
  226. - PDF输入: 文件名_page_001_ocr.png
  227. - 图片输入(单页): 文件名_ocr.png
  228. Args:
  229. results: 处理结果
  230. output_dir: 输出目录
  231. doc_name: 文档名称
  232. is_pdf: 是否为 PDF 输入
  233. Returns:
  234. 保存的图片路径列表
  235. """
  236. ocr_paths = []
  237. total_pages = len(results.get('pages', []))
  238. for page in results.get('pages', []):
  239. page_idx = page.get('page_idx', 0)
  240. processed_image = page.get('original_image')
  241. if processed_image is None:
  242. processed_image = page.get('processed_image')
  243. if processed_image is None:
  244. logger.warning(f"Page {page_idx}: No image data found for OCR visualization")
  245. continue
  246. if isinstance(processed_image, np.ndarray):
  247. image = Image.fromarray(processed_image).convert('RGB')
  248. elif isinstance(processed_image, Image.Image):
  249. image = processed_image.convert('RGB')
  250. else:
  251. continue
  252. draw = ImageDraw.Draw(image)
  253. font = VisualizationUtils._get_font(10)
  254. for element in page.get('elements', []):
  255. content = element.get('content', {})
  256. # OCR 文本框
  257. ocr_details = content.get('ocr_details', [])
  258. for ocr_item in ocr_details:
  259. ocr_bbox = ocr_item.get('bbox', [])
  260. if ocr_bbox:
  261. VisualizationUtils._draw_polygon(
  262. draw, ocr_bbox, VisualizationUtils.OCR_BOX_COLOR, width=1
  263. )
  264. # 表格单元格
  265. cells = content.get('cells', [])
  266. for cell in cells:
  267. cell_bbox = cell.get('bbox', [])
  268. if cell_bbox and len(cell_bbox) >= 4:
  269. x0, y0, x1, y1 = map(int, cell_bbox[:4])
  270. draw.rectangle(
  271. [x0, y0, x1, y1],
  272. outline=VisualizationUtils.CELL_BOX_COLOR,
  273. width=2
  274. )
  275. cell_text = cell.get('text', '')[:10]
  276. if cell_text:
  277. draw.text(
  278. (x0 + 2, y0 + 2),
  279. cell_text,
  280. fill=VisualizationUtils.CELL_BOX_COLOR,
  281. font=font
  282. )
  283. # OCR 框
  284. ocr_boxes = content.get('ocr_boxes', [])
  285. for ocr_box in ocr_boxes:
  286. bbox = ocr_box.get('bbox', [])
  287. if bbox:
  288. VisualizationUtils._draw_polygon(
  289. draw, bbox, VisualizationUtils.OCR_BOX_COLOR, width=1
  290. )
  291. # 绘制丢弃元素的 OCR 框
  292. for element in page.get('discarded_blocks', []):
  293. bbox = element.get('bbox', [0, 0, 0, 0])
  294. content = element.get('content', {})
  295. if len(bbox) >= 4:
  296. x0, y0, x1, y1 = map(int, bbox[:4])
  297. draw.rectangle(
  298. [x0, y0, x1, y1],
  299. outline=VisualizationUtils.DISCARD_COLOR,
  300. width=1
  301. )
  302. ocr_details = content.get('ocr_details', [])
  303. for ocr_item in ocr_details:
  304. ocr_bbox = ocr_item.get('bbox', [])
  305. if ocr_bbox:
  306. VisualizationUtils._draw_polygon(
  307. draw, ocr_bbox, VisualizationUtils.DISCARD_COLOR, width=1
  308. )
  309. # 根据输入类型决定命名
  310. if is_pdf or total_pages > 1:
  311. ocr_path = output_dir / f"{doc_name}_page_{page_idx + 1:03d}_ocr.png"
  312. else:
  313. ocr_path = output_dir / f"{doc_name}_ocr.png"
  314. image.save(ocr_path)
  315. ocr_paths.append(str(ocr_path))
  316. logger.info(f"🖼️ OCR image saved: {ocr_path}")
  317. return ocr_paths
  318. @staticmethod
  319. def _draw_polygon(
  320. draw: ImageDraw.Draw,
  321. bbox: List,
  322. color: Tuple[int, int, int],
  323. width: int = 1
  324. ):
  325. """
  326. 绘制多边形或矩形
  327. Args:
  328. draw: ImageDraw 对象
  329. bbox: 坐标(4点多边形或矩形)
  330. color: 颜色
  331. width: 线宽
  332. """
  333. if isinstance(bbox[0], (list, tuple)):
  334. points = [(int(p[0]), int(p[1])) for p in bbox]
  335. points.append(points[0])
  336. draw.line(points, fill=color, width=width)
  337. elif len(bbox) >= 4:
  338. x0, y0, x1, y1 = map(int, bbox[:4])
  339. draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
  340. @staticmethod
  341. def _get_font(size: int) -> ImageFont.FreeTypeFont:
  342. """
  343. 获取字体
  344. Args:
  345. size: 字体大小
  346. Returns:
  347. 字体对象
  348. """
  349. font_paths = [
  350. "/System/Library/Fonts/Helvetica.ttc",
  351. "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
  352. "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
  353. ]
  354. for font_path in font_paths:
  355. try:
  356. return ImageFont.truetype(font_path, size)
  357. except:
  358. continue
  359. return ImageFont.load_default()