|
@@ -0,0 +1,723 @@
|
|
|
|
|
+"""
|
|
|
|
|
+增强版输出格式化器 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"<!-- bbox: {bbox} -->"
|
|
|
|
|
+
|
|
|
|
|
+ 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("<details>")
|
|
|
|
|
+ md_lines.append("<summary>单元格坐标信息</summary>")
|
|
|
|
|
+ md_lines.append("")
|
|
|
|
|
+ md_lines.append("```json")
|
|
|
|
|
+ md_lines.append(json.dumps(cells, ensure_ascii=False, indent=2))
|
|
|
|
|
+ md_lines.append("```")
|
|
|
|
|
+ md_lines.append("</details>")
|
|
|
|
|
+ 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"""<!DOCTYPE html>
|
|
|
|
|
+<html lang="zh-CN">
|
|
|
|
|
+<head>
|
|
|
|
|
+ <meta charset="UTF-8">
|
|
|
|
|
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
|
|
|
|
|
+ <title>{doc_name} - Table {table_idx}</title>
|
|
|
|
|
+ <style>
|
|
|
|
|
+ body {{
|
|
|
|
|
+ font-family: Arial, "Microsoft YaHei", sans-serif;
|
|
|
|
|
+ margin: 20px;
|
|
|
|
|
+ background-color: #f5f5f5;
|
|
|
|
|
+ }}
|
|
|
|
|
+ .container {{
|
|
|
|
|
+ max-width: 1400px;
|
|
|
|
|
+ margin: 0 auto;
|
|
|
|
|
+ background-color: white;
|
|
|
|
|
+ padding: 20px;
|
|
|
|
|
+ box-shadow: 0 0 10px rgba(0,0,0,0.1);
|
|
|
|
|
+ border-radius: 8px;
|
|
|
|
|
+ }}
|
|
|
|
|
+ .meta {{
|
|
|
|
|
+ color: #666;
|
|
|
|
|
+ font-size: 0.9em;
|
|
|
|
|
+ margin-bottom: 20px;
|
|
|
|
|
+ padding-bottom: 10px;
|
|
|
|
|
+ border-bottom: 1px solid #ddd;
|
|
|
|
|
+ }}
|
|
|
|
|
+ table {{
|
|
|
|
|
+ border-collapse: collapse;
|
|
|
|
|
+ width: 100%;
|
|
|
|
|
+ margin: 20px 0;
|
|
|
|
|
+ }}
|
|
|
|
|
+ th, td {{
|
|
|
|
|
+ border: 1px solid #ddd;
|
|
|
|
|
+ padding: 8px 12px;
|
|
|
|
|
+ text-align: left;
|
|
|
|
|
+ }}
|
|
|
|
|
+ th {{
|
|
|
|
|
+ background-color: #f2f2f2;
|
|
|
|
|
+ font-weight: bold;
|
|
|
|
|
+ }}
|
|
|
|
|
+ tr:hover {{
|
|
|
|
|
+ background-color: #f9f9f9;
|
|
|
|
|
+ }}
|
|
|
|
|
+ td[data-bbox], th[data-bbox] {{
|
|
|
|
|
+ position: relative;
|
|
|
|
|
+ }}
|
|
|
|
|
+ td[data-bbox]:hover::after, th[data-bbox]:hover::after {{
|
|
|
|
|
+ content: attr(data-bbox);
|
|
|
|
|
+ position: absolute;
|
|
|
|
|
+ bottom: 100%;
|
|
|
|
|
+ left: 0;
|
|
|
|
|
+ background: #333;
|
|
|
|
|
+ color: white;
|
|
|
|
|
+ padding: 2px 6px;
|
|
|
|
|
+ font-size: 10px;
|
|
|
|
|
+ border-radius: 3px;
|
|
|
|
|
+ white-space: nowrap;
|
|
|
|
|
+ z-index: 100;
|
|
|
|
|
+ }}
|
|
|
|
|
+ .cells-info {{
|
|
|
|
|
+ margin-top: 30px;
|
|
|
|
|
+ padding: 15px;
|
|
|
|
|
+ background-color: #f8f9fa;
|
|
|
|
|
+ border-radius: 5px;
|
|
|
|
|
+ }}
|
|
|
|
|
+ .cells-info summary {{
|
|
|
|
|
+ cursor: pointer;
|
|
|
|
|
+ font-weight: bold;
|
|
|
|
|
+ color: #333;
|
|
|
|
|
+ }}
|
|
|
|
|
+ .cells-info pre {{
|
|
|
|
|
+ background-color: #2d2d2d;
|
|
|
|
|
+ color: #f8f8f2;
|
|
|
|
|
+ padding: 15px;
|
|
|
|
|
+ border-radius: 5px;
|
|
|
|
|
+ overflow-x: auto;
|
|
|
|
|
+ font-size: 12px;
|
|
|
|
|
+ }}
|
|
|
|
|
+ </style>
|
|
|
|
|
+</head>
|
|
|
|
|
+<body>
|
|
|
|
|
+ <div class="container">
|
|
|
|
|
+ <div class="meta">
|
|
|
|
|
+ <p><strong>Document:</strong> {doc_name}</p>
|
|
|
|
|
+ <p><strong>Page:</strong> {page_idx + 1}</p>
|
|
|
|
|
+ <p><strong>Table:</strong> {table_idx}</p>
|
|
|
|
|
+ <p><strong>Cells with coordinates:</strong> {len(cells)}</p>
|
|
|
|
|
+ </div>
|
|
|
|
|
+
|
|
|
|
|
+ {table_html}
|
|
|
|
|
+
|
|
|
|
|
+ <div class="cells-info">
|
|
|
|
|
+ <details>
|
|
|
|
|
+ <summary>📍 单元格坐标数据 (JSON)</summary>
|
|
|
|
|
+ <pre>{cells_json}</pre>
|
|
|
|
|
+ </details>
|
|
|
|
|
+ </div>
|
|
|
|
|
+ </div>
|
|
|
|
|
+</body>
|
|
|
|
|
+</html>"""
|
|
|
|
|
+
|
|
|
|
|
+ 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": "<table><tr><td>日期</td><td>金额</td></tr></table>",
|
|
|
|
|
+ "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)
|
|
|
|
|
+
|