MinerU/zhch/universal_doc_parser/utils/visualization_utils.py

"""
可视化工具模块

提供文档处理结果的可视化功能：
- 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()