| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370 |
- """
- 可视化工具模块
- 提供文档处理结果的可视化功能:
- - Layout 布局可视化
- - OCR 结果可视化
- - 图片元素保存
- """
- from pathlib import Path
- from typing import Dict, Any, List, Tuple
- import numpy as np
- from PIL import Image, ImageDraw, ImageFont
- import cv2
- from loguru import logger
- class VisualizationUtils:
- """可视化工具类"""
-
- # 颜色映射(与 MinerU 保持一致)
- COLOR_MAP = {
- 'title': (102, 102, 255), # 蓝色
- 'text': (153, 0, 76), # 深红
- 'image': (153, 255, 51), # 绿色
- 'image_body': (153, 255, 51),
- 'image_caption': (102, 178, 255),
- 'image_footnote': (255, 178, 102),
- 'table': (204, 204, 0), # 黄色
- 'table_body': (204, 204, 0),
- 'table_caption': (255, 255, 102),
- 'table_footnote': (229, 255, 204),
- 'interline_equation': (0, 255, 0), # 亮绿
- 'inline_equation': (0, 200, 0),
- 'list': (40, 169, 92),
- 'code': (102, 0, 204), # 紫色
- 'header': (128, 128, 128), # 灰色
- 'footer': (128, 128, 128),
- 'ref_text': (180, 180, 180),
- 'ocr_text': (153, 0, 76),
- 'error': (255, 0, 0), # 红色
- }
-
- # OCR 框颜色
- OCR_BOX_COLOR = (0, 255, 0) # 绿色
- CELL_BOX_COLOR = (255, 165, 0) # 橙色
- DISCARD_COLOR = (128, 128, 128) # 灰色
-
- @staticmethod
- def save_image_elements(
- results: Dict[str, Any],
- images_dir: Path,
- doc_name: str
- ) -> List[str]:
- """
- 保存图片元素
-
- Args:
- results: 处理结果
- images_dir: 图片输出目录
- doc_name: 文档名称
-
- Returns:
- 保存的图片路径列表
- """
- saved_paths = []
- image_count = 0
-
- for page in results.get('pages', []):
- page_idx = page.get('page_idx', 0)
-
- for element in page.get('elements', []):
- if element.get('type') in ['image', 'image_body', 'figure']:
- content = element.get('content', {})
- image_data = content.get('image_data')
-
- if image_data is not None:
- image_count += 1
- image_filename = f"{doc_name}_page_{page_idx + 1}_image_{image_count}.png"
- image_path = images_dir / image_filename
-
- try:
- if isinstance(image_data, np.ndarray):
- cv2.imwrite(str(image_path), image_data)
- else:
- Image.fromarray(image_data).save(image_path)
-
- # 更新路径(只保存文件名)
- content['image_path'] = image_filename
- content.pop('image_data', None)
-
- saved_paths.append(str(image_path))
- logger.debug(f"🖼️ Image saved: {image_path}")
- except Exception as e:
- logger.warning(f"Failed to save image: {e}")
-
- if image_count > 0:
- logger.info(f"🖼️ {image_count} images saved to: {images_dir}")
-
- return saved_paths
-
- @staticmethod
- def save_layout_images(
- results: Dict[str, Any],
- output_dir: Path,
- doc_name: str,
- draw_type_label: bool = True,
- draw_bbox_number: bool = True
- ) -> List[str]:
- """
- 保存 Layout 可视化图片
-
- Args:
- results: 处理结果
- output_dir: 输出目录
- doc_name: 文档名称
- draw_type_label: 是否绘制类型标签
- draw_bbox_number: 是否绘制序号
-
- Returns:
- 保存的图片路径列表
- """
- layout_paths = []
-
- for page in results.get('pages', []):
- page_idx = page.get('page_idx', 0)
- processed_image = page.get('original_image')
- if processed_image is None:
- processed_image = page.get('processed_image')
-
- if processed_image is None:
- logger.warning(f"Page {page_idx}: No image data found for layout visualization")
- continue
-
- if isinstance(processed_image, np.ndarray):
- image = Image.fromarray(processed_image).convert('RGB')
- elif isinstance(processed_image, Image.Image):
- image = processed_image.convert('RGB')
- else:
- continue
-
- draw = ImageDraw.Draw(image, 'RGBA')
- font = VisualizationUtils._get_font(14)
-
- # 绘制普通元素
- for idx, element in enumerate(page.get('elements', []), 1):
- elem_type = element.get('type', '')
- bbox = element.get('bbox', [0, 0, 0, 0])
-
- if len(bbox) < 4:
- continue
-
- x0, y0, x1, y1 = map(int, bbox[:4])
- color = VisualizationUtils.COLOR_MAP.get(elem_type, (255, 0, 0))
-
- # 半透明填充
- overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
- overlay_draw = ImageDraw.Draw(overlay)
- overlay_draw.rectangle([x0, y0, x1, y1], fill=(*color, 50))
- image = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
- draw = ImageDraw.Draw(image)
-
- # 边框
- draw.rectangle([x0, y0, x1, y1], outline=color, width=2)
-
- # 类型标签
- if draw_type_label:
- label = elem_type.replace('_', ' ').title()
- bbox_label = draw.textbbox((x0 + 2, y0 + 2), label, font=font)
- draw.rectangle(bbox_label, fill=color)
- draw.text((x0 + 2, y0 + 2), label, fill='white', font=font)
-
- # 序号
- if draw_bbox_number:
- number_text = str(idx)
- bbox_number = draw.textbbox((x1 - 25, y0 + 2), number_text, font=font)
- draw.rectangle(bbox_number, fill=(255, 0, 0))
- draw.text((x1 - 25, y0 + 2), number_text, fill='white', font=font)
-
- # 绘制丢弃元素(灰色样式)
- for idx, element in enumerate(page.get('discarded_blocks', []), 1):
- original_category = element.get('original_category', 'unknown')
- bbox = element.get('bbox', [0, 0, 0, 0])
-
- if len(bbox) < 4:
- continue
-
- x0, y0, x1, y1 = map(int, bbox[:4])
-
- # 半透明填充
- overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
- overlay_draw = ImageDraw.Draw(overlay)
- overlay_draw.rectangle([x0, y0, x1, y1], fill=(*VisualizationUtils.DISCARD_COLOR, 30))
- image = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
- draw = ImageDraw.Draw(image)
-
- # 灰色边框
- draw.rectangle([x0, y0, x1, y1], outline=VisualizationUtils.DISCARD_COLOR, width=1)
-
- # 类型标签
- if draw_type_label:
- label = f"D:{original_category}"
- bbox_label = draw.textbbox((x0 + 2, y0 + 2), label, font=font)
- draw.rectangle(bbox_label, fill=VisualizationUtils.DISCARD_COLOR)
- draw.text((x0 + 2, y0 + 2), label, fill='white', font=font)
-
- layout_path = output_dir / f"{doc_name}_page_{page_idx + 1}_layout.png"
- image.save(layout_path)
- layout_paths.append(str(layout_path))
- logger.info(f"🖼️ Layout image saved: {layout_path}")
-
- return layout_paths
-
- @staticmethod
- def save_ocr_images(
- results: Dict[str, Any],
- output_dir: Path,
- doc_name: str
- ) -> List[str]:
- """
- 保存 OCR 可视化图片
-
- Args:
- results: 处理结果
- output_dir: 输出目录
- doc_name: 文档名称
-
- Returns:
- 保存的图片路径列表
- """
- ocr_paths = []
-
- for page in results.get('pages', []):
- page_idx = page.get('page_idx', 0)
- processed_image = page.get('original_image')
- if processed_image is None:
- processed_image = page.get('processed_image')
-
- if processed_image is None:
- logger.warning(f"Page {page_idx}: No image data found for OCR visualization")
- continue
-
- if isinstance(processed_image, np.ndarray):
- image = Image.fromarray(processed_image).convert('RGB')
- elif isinstance(processed_image, Image.Image):
- image = processed_image.convert('RGB')
- else:
- continue
-
- draw = ImageDraw.Draw(image)
- font = VisualizationUtils._get_font(10)
-
- for element in page.get('elements', []):
- content = element.get('content', {})
-
- # OCR 文本框
- ocr_details = content.get('ocr_details', [])
- for ocr_item in ocr_details:
- ocr_bbox = ocr_item.get('bbox', [])
- if ocr_bbox:
- VisualizationUtils._draw_polygon(
- draw, ocr_bbox, VisualizationUtils.OCR_BOX_COLOR, width=1
- )
-
- # 表格单元格
- cells = content.get('cells', [])
- for cell in cells:
- cell_bbox = cell.get('bbox', [])
- if cell_bbox and len(cell_bbox) >= 4:
- x0, y0, x1, y1 = map(int, cell_bbox[:4])
- draw.rectangle(
- [x0, y0, x1, y1],
- outline=VisualizationUtils.CELL_BOX_COLOR,
- width=2
- )
-
- cell_text = cell.get('text', '')[:10]
- if cell_text:
- draw.text(
- (x0 + 2, y0 + 2),
- cell_text,
- fill=VisualizationUtils.CELL_BOX_COLOR,
- font=font
- )
-
- # OCR 框
- ocr_boxes = content.get('ocr_boxes', [])
- for ocr_box in ocr_boxes:
- bbox = ocr_box.get('bbox', [])
- if bbox:
- VisualizationUtils._draw_polygon(
- draw, bbox, VisualizationUtils.OCR_BOX_COLOR, width=1
- )
-
- # 绘制丢弃元素的 OCR 框
- for element in page.get('discarded_blocks', []):
- bbox = element.get('bbox', [0, 0, 0, 0])
- content = element.get('content', {})
-
- if len(bbox) >= 4:
- x0, y0, x1, y1 = map(int, bbox[:4])
- draw.rectangle(
- [x0, y0, x1, y1],
- outline=VisualizationUtils.DISCARD_COLOR,
- width=1
- )
-
- ocr_details = content.get('ocr_details', [])
- for ocr_item in ocr_details:
- ocr_bbox = ocr_item.get('bbox', [])
- if ocr_bbox:
- VisualizationUtils._draw_polygon(
- draw, ocr_bbox, VisualizationUtils.DISCARD_COLOR, width=1
- )
-
- ocr_path = output_dir / f"{doc_name}_page_{page_idx + 1}_ocr.png"
- image.save(ocr_path)
- ocr_paths.append(str(ocr_path))
- logger.info(f"🖼️ OCR image saved: {ocr_path}")
-
- return ocr_paths
-
- @staticmethod
- def _draw_polygon(
- draw: ImageDraw.Draw,
- bbox: List,
- color: Tuple[int, int, int],
- width: int = 1
- ):
- """
- 绘制多边形或矩形
-
- Args:
- draw: ImageDraw 对象
- bbox: 坐标(4点多边形或矩形)
- color: 颜色
- width: 线宽
- """
- if isinstance(bbox[0], (list, tuple)):
- points = [(int(p[0]), int(p[1])) for p in bbox]
- points.append(points[0])
- draw.line(points, fill=color, width=width)
- elif len(bbox) >= 4:
- x0, y0, x1, y1 = map(int, bbox[:4])
- draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
-
- @staticmethod
- def _get_font(size: int) -> ImageFont.FreeTypeFont:
- """
- 获取字体
-
- Args:
- size: 字体大小
-
- Returns:
- 字体对象
- """
- font_paths = [
- "/System/Library/Fonts/Helvetica.ttc",
- "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
- "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
- ]
-
- for font_path in font_paths:
- try:
- return ImageFont.truetype(font_path, size)
- except:
- continue
-
- return ImageFont.load_default()
|