""" 增强版输出格式化器 v2 支持: 1. MinerU 标准 middle.json 格式 2. Markdown 输出 3. Debug 模式:layout 图片、OCR 图片、单元格坐标图片 4. 表格 HTML 输出(带坐标信息) """ import json import os import sys from pathlib import Path from typing import Dict, Any, List, Optional, Union from loguru import logger import numpy as np from PIL import Image, ImageDraw, ImageFont import cv2 # 导入 MinerU 组件 mineru_path = Path(__file__).parents[3] if str(mineru_path) not in sys.path: sys.path.insert(0, str(mineru_path)) try: from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make from mineru.utils.enum_class import MakeMode, BlockType, ContentType MINERU_AVAILABLE = True except ImportError as e: logger.warning(f"MinerU components not available: {e}") MINERU_AVAILABLE = False class OutputFormatterV2: """增强版输出格式化器""" # 颜色映射(与 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) # 橙色 def __init__(self, output_dir: str): """ 初始化格式化器 Args: output_dir: 输出目录 """ self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) def save_results( self, results: Dict[str, Any], output_config: Dict[str, Any] ) -> Dict[str, str]: """ 保存处理结果 Args: results: 处理结果 output_config: 输出配置 Returns: 输出文件路径字典 """ output_paths = {} # 创建文档输出目录 doc_name = Path(results['document_path']).stem doc_output_dir = self.output_dir / doc_name doc_output_dir.mkdir(parents=True, exist_ok=True) # 1. 转换为 MinerU 标准 middle.json middle_json = self._convert_to_middle_json(results) # 2. 保存 middle.json if output_config.get('save_json', True): json_path = doc_output_dir / f"{doc_name}_middle.json" with open(json_path, 'w', encoding='utf-8') as f: json.dump(middle_json, f, ensure_ascii=False, indent=2) output_paths['middle_json'] = str(json_path) logger.info(f"📄 Middle JSON saved: {json_path}") # 3. 保存增强版 JSON(包含单元格坐标) enhanced_json_path = doc_output_dir / f"{doc_name}_enhanced.json" with open(enhanced_json_path, 'w', encoding='utf-8') as f: json.dump(results, f, ensure_ascii=False, indent=2, default=self._json_serializer) output_paths['enhanced_json'] = str(enhanced_json_path) logger.info(f"📄 Enhanced JSON saved: {enhanced_json_path}") # 4. 保存 Markdown if output_config.get('save_markdown', True): md_path = self._save_markdown(results, doc_output_dir, doc_name) output_paths['markdown'] = str(md_path) # 5. 保存表格 HTML if output_config.get('save_html', True): html_dir = self._save_table_htmls(results, doc_output_dir, doc_name) output_paths['table_htmls'] = str(html_dir) # 6. Debug 模式:保存可视化图片 if output_config.get('save_layout_image', False): layout_paths = self._save_layout_images( results, doc_output_dir, doc_name, draw_type_label=output_config.get('draw_type_label', True), draw_bbox_number=output_config.get('draw_bbox_number', True) ) output_paths['layout_images'] = layout_paths if output_config.get('save_ocr_image', False): ocr_paths = self._save_ocr_images(results, doc_output_dir, doc_name) output_paths['ocr_images'] = ocr_paths logger.info(f"✅ All results saved to: {doc_output_dir}") return output_paths def _json_serializer(self, obj): """JSON 序列化辅助函数""" if isinstance(obj, np.ndarray): return obj.tolist() if isinstance(obj, np.integer): return int(obj) if isinstance(obj, np.floating): return float(obj) if isinstance(obj, Path): return str(obj) return str(obj) def _convert_to_middle_json(self, results: Dict[str, Any]) -> Dict[str, Any]: """转换为 MinerU 标准 middle.json 格式""" middle_json = { "pdf_info": [], "_backend": "vlm", "_scene": results.get('scene', 'unknown'), "_version_name": "2.5.0" } for page in results.get('pages', []): page_info = { 'page_idx': page['page_idx'], 'page_size': list(page.get('image_shape', [0, 0])[:2][::-1]), 'angle': page.get('angle', 0), 'para_blocks': [], 'discarded_blocks': [] } for element in page.get('elements', []): block = self._element_to_block(element, page_info['page_size']) if block: elem_type = element.get('type', '') if elem_type in ['header', 'footer', 'page_number', 'aside_text']: page_info['discarded_blocks'].append(block) else: page_info['para_blocks'].append(block) middle_json['pdf_info'].append(page_info) return middle_json def _element_to_block( self, element: Dict[str, Any], page_size: List[int] ) -> Dict[str, Any]: """将元素转换为 MinerU block 格式""" elem_type = element.get('type', '') bbox = element.get('bbox', [0, 0, 0, 0]) content = element.get('content', {}) block = { 'type': elem_type, 'bbox': bbox, 'angle': element.get('angle', 0), 'lines': [] } # 文本类型 if elem_type in ['text', 'title', 'ref_text', 'header', 'footer', 'ocr_text']: text = content.get('text', '') if isinstance(content, dict) else str(content) if text: block['lines'] = [{ 'bbox': bbox, 'spans': [{ 'bbox': bbox, 'type': 'text', 'content': text }] }] # 表格类型 elif elem_type in ['table', 'table_body']: table_html = content.get('html', '') cells = content.get('cells', []) block['blocks'] = [{ 'type': 'table_body', 'bbox': bbox, 'angle': 0, 'lines': [{ 'bbox': bbox, 'spans': [{ 'bbox': bbox, 'type': 'table', 'html': table_html, 'cells': cells # 增强:包含单元格坐标 }] }] }] # 公式类型 elif elem_type in ['interline_equation', 'inline_equation', 'equation']: latex = content.get('latex', '') block['lines'] = [{ 'bbox': bbox, 'spans': [{ 'bbox': bbox, 'type': 'interline_equation' if 'interline' in elem_type else 'inline_equation', 'content': latex }] }] return block def _save_markdown( self, results: Dict[str, Any], output_dir: Path, doc_name: str ) -> Path: """保存 Markdown 文件""" md_path = output_dir / f"{doc_name}.md" md_lines = [ f"---", f"scene: {results.get('scene', 'unknown')}", f"document: {results.get('document_path', '')}", f"pages: {len(results.get('pages', []))}", f"---", "", ] for page in results.get('pages', []): page_idx = page.get('page_idx', 0) md_lines.append(f"\n## Page {page_idx + 1}\n") for element in page.get('elements', []): elem_type = element.get('type', '') content = element.get('content', {}) bbox = element.get('bbox', []) # 添加 bbox 注释 bbox_comment = f"" if elem_type == 'title': text = content.get('text', '') if isinstance(content, dict) else str(content) level = element.get('level', 1) md_lines.append(f"{'#' * min(level, 6)} {text}") md_lines.append(bbox_comment) md_lines.append("") elif elem_type in ['text', 'ocr_text', 'ref_text']: text = content.get('text', '') if isinstance(content, dict) else str(content) if text: md_lines.append(text) md_lines.append(bbox_comment) md_lines.append("") elif elem_type in ['table', 'table_body']: html = content.get('html', '') if html: # 转换 HTML 表格为 Markdown md_table = self._html_to_markdown_table(html) md_lines.append(md_table) md_lines.append(bbox_comment) # 添加单元格坐标信息 cells = content.get('cells', []) if cells: md_lines.append("") md_lines.append("
") md_lines.append("单元格坐标信息") md_lines.append("") md_lines.append("```json") md_lines.append(json.dumps(cells, ensure_ascii=False, indent=2)) md_lines.append("```") md_lines.append("
") md_lines.append("") elif elem_type in ['interline_equation', 'inline_equation', 'equation']: latex = content.get('latex', '') if latex: md_lines.append(f"$$\n{latex}\n$$") md_lines.append(bbox_comment) md_lines.append("") with open(md_path, 'w', encoding='utf-8') as f: f.write('\n'.join(md_lines)) logger.info(f"📝 Markdown saved: {md_path}") return md_path def _html_to_markdown_table(self, html: str) -> str: """将 HTML 表格转换为 Markdown 格式""" try: from bs4 import BeautifulSoup soup = BeautifulSoup(html, 'html.parser') rows = soup.find_all('tr') if not rows: return html md_rows = [] for row_idx, row in enumerate(rows): cells = row.find_all(['td', 'th']) cell_texts = [cell.get_text(strip=True) for cell in cells] md_row = '| ' + ' | '.join(cell_texts) + ' |' md_rows.append(md_row) # 添加表头分隔符 if row_idx == 0: separator = '| ' + ' | '.join(['---'] * len(cells)) + ' |' md_rows.append(separator) return '\n'.join(md_rows) except Exception as e: logger.warning(f"HTML to Markdown conversion failed: {e}") return html def _save_table_htmls( self, results: Dict[str, Any], output_dir: Path, doc_name: str ) -> Path: """保存表格 HTML 文件(带坐标信息)""" tables_dir = output_dir / 'tables' tables_dir.mkdir(exist_ok=True) table_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 ['table', 'table_body']: table_count += 1 content = element.get('content', {}) html = content.get('html', '') cells = content.get('cells', []) if html: # 生成带样式的 HTML full_html = self._generate_table_html_with_styles( html, cells, doc_name, page_idx, table_count ) html_path = tables_dir / f"{doc_name}_table_{table_count}_page_{page_idx + 1}.html" with open(html_path, 'w', encoding='utf-8') as f: f.write(full_html) if table_count > 0: logger.info(f"📊 {table_count} tables saved to: {tables_dir}") return tables_dir def _generate_table_html_with_styles( self, table_html: str, cells: List[Dict], doc_name: str, page_idx: int, table_idx: int ) -> str: """生成带样式和坐标信息的完整 HTML""" cells_json = json.dumps(cells, ensure_ascii=False, indent=2) if cells else "[]" return f""" {doc_name} - Table {table_idx}

Document: {doc_name}

Page: {page_idx + 1}

Table: {table_idx}

Cells with coordinates: {len(cells)}

{table_html}
📍 单元格坐标数据 (JSON)
{cells_json}
""" def _save_layout_images( self, results: Dict[str, Any], output_dir: Path, doc_name: str, draw_type_label: bool = True, draw_bbox_number: bool = True ) -> List[str]: """保存 Layout 可视化图片""" layout_paths = [] for page in results.get('pages', []): page_idx = page.get('page_idx', 0) processed_image = page.get('processed_image') if processed_image is None: continue # 转换为 PIL Image 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 = self._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 = self.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) # 保存图片 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 def _save_ocr_images( self, results: Dict[str, Any], output_dir: Path, doc_name: str ) -> List[str]: """保存 OCR 可视化图片(显示文本框和单元格坐标)""" ocr_paths = [] for page in results.get('pages', []): page_idx = page.get('page_idx', 0) processed_image = page.get('processed_image') if processed_image is None: continue # 转换为 PIL Image 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 = self._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: self._draw_polygon(draw, ocr_bbox, self.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=self.CELL_BOX_COLOR, width=2) # 标注单元格文本 cell_text = cell.get('text', '')[:10] if cell_text: draw.text((x0 + 2, y0 + 2), cell_text, fill=self.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: self._draw_polygon(draw, bbox, self.OCR_BOX_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 def _draw_polygon( self, draw: ImageDraw.Draw, bbox: List, color: tuple, width: int = 1 ): """绘制多边形或矩形""" # 4点坐标格式 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]] 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) # 4值坐标格式 [x1, y1, x2, y2] elif len(bbox) >= 4: x0, y0, x1, y1 = map(int, bbox[:4]) draw.rectangle([x0, y0, x1, y1], outline=color, width=width) def _get_font(self, size: int) -> ImageFont.FreeTypeFont: """获取字体""" 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() if __name__ == "__main__": # 测试代码 sample_results = { "document_path": "/path/to/sample.pdf", "scene": "bank_statement", "pages": [ { "page_idx": 0, "image_shape": [1654, 2338, 3], "elements": [ { "type": "title", "bbox": [100, 50, 800, 100], "content": {"text": "银行流水"}, "confidence": 0.98 }, { "type": "table", "bbox": [100, 200, 800, 600], "content": { "html": "
日期金额
", "cells": [ {"text": "日期", "bbox": [100, 200, 200, 250], "row": 1, "col": 1}, {"text": "金额", "bbox": [200, 200, 300, 250], "row": 1, "col": 2} ] } } ] } ] } formatter = OutputFormatterV2("./test_output_v2") output_files = formatter.save_results( sample_results, { "save_json": True, "save_markdown": True, "save_html": True, "save_layout_image": False, "save_ocr_image": False } ) print("Generated files:", output_files)