| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723 |
- """
- 增强版输出格式化器 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)
|