|
|
@@ -1,16 +1,20 @@
|
|
|
"""
|
|
|
-增强版输出格式化器 v2
|
|
|
+统一输出格式化器
|
|
|
+严格遵循 MinerU mineru_vllm_results_cell_bbox 格式
|
|
|
+
|
|
|
支持:
|
|
|
-1. MinerU 标准 middle.json 格式
|
|
|
-2. Markdown 输出
|
|
|
-3. Debug 模式:layout 图片、OCR 图片、单元格坐标图片
|
|
|
-4. 表格 HTML 输出(带坐标信息)
|
|
|
+1. MinerU 标准 middle.json 格式(用于 union_make 生成 Markdown)
|
|
|
+2. mineru_vllm_results_cell_bbox 格式(每页独立 JSON)
|
|
|
+3. Markdown 输出(复用 MinerU union_make)
|
|
|
+4. Debug 模式:layout 图片、OCR 图片
|
|
|
+5. 表格 HTML 输出(带坐标信息)
|
|
|
"""
|
|
|
import json
|
|
|
import os
|
|
|
import sys
|
|
|
+import hashlib
|
|
|
from pathlib import Path
|
|
|
-from typing import Dict, Any, List, Optional, Union
|
|
|
+from typing import Dict, Any, List, Optional, Union, Tuple
|
|
|
from loguru import logger
|
|
|
import numpy as np
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
@@ -28,10 +32,23 @@ try:
|
|
|
except ImportError as e:
|
|
|
logger.warning(f"MinerU components not available: {e}")
|
|
|
MINERU_AVAILABLE = False
|
|
|
+
|
|
|
+ # 占位符定义
|
|
|
+ class MakeMode:
|
|
|
+ MM_MD = 'mm_md'
|
|
|
+ NLP_MD = 'nlp_md'
|
|
|
|
|
|
|
|
|
class OutputFormatterV2:
|
|
|
- """增强版输出格式化器"""
|
|
|
+ """
|
|
|
+ 统一输出格式化器
|
|
|
+
|
|
|
+ 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式:
|
|
|
+ - middle.json: MinerU 标准格式,用于生成 Markdown
|
|
|
+ - page_xxx.json: 每页独立的 JSON,包含 table_cells
|
|
|
+ - Markdown: 带 bbox 注释
|
|
|
+ - 表格: HTML 格式,带 data-bbox 属性
|
|
|
+ """
|
|
|
|
|
|
# 颜色映射(与 MinerU 保持一致)
|
|
|
COLOR_MAP = {
|
|
|
@@ -74,7 +91,7 @@ class OutputFormatterV2:
|
|
|
self,
|
|
|
results: Dict[str, Any],
|
|
|
output_config: Dict[str, Any]
|
|
|
- ) -> Dict[str, str]:
|
|
|
+ ) -> Dict[str, Any]:
|
|
|
"""
|
|
|
保存处理结果
|
|
|
|
|
|
@@ -85,17 +102,29 @@ class OutputFormatterV2:
|
|
|
Returns:
|
|
|
输出文件路径字典
|
|
|
"""
|
|
|
- output_paths = {}
|
|
|
+ output_paths = {
|
|
|
+ 'images': [],
|
|
|
+ 'json_pages': [],
|
|
|
+ }
|
|
|
|
|
|
# 创建文档输出目录
|
|
|
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
|
|
|
+ # 创建 images 子目录
|
|
|
+ images_dir = doc_output_dir / 'images'
|
|
|
+ images_dir.mkdir(exist_ok=True)
|
|
|
+
|
|
|
+ # 1. 首先保存图片元素(设置 image_path)
|
|
|
+ image_paths = self._save_image_elements(results, images_dir, doc_name)
|
|
|
+ if image_paths:
|
|
|
+ output_paths['images'] = image_paths
|
|
|
+
|
|
|
+ # 2. 转换为 MinerU middle.json 格式(用于 union_make)
|
|
|
middle_json = self._convert_to_middle_json(results)
|
|
|
|
|
|
- # 2. 保存 middle.json
|
|
|
+ # 3. 保存 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:
|
|
|
@@ -103,24 +132,22 @@ class OutputFormatterV2:
|
|
|
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. 保存每页独立的 mineru_vllm_results_cell_bbox 格式 JSON
|
|
|
+ if output_config.get('save_page_json', True):
|
|
|
+ page_json_paths = self._save_page_jsons(results, doc_output_dir, doc_name)
|
|
|
+ output_paths['json_pages'] = page_json_paths
|
|
|
|
|
|
- # 4. 保存 Markdown
|
|
|
+ # 5. 保存 Markdown
|
|
|
if output_config.get('save_markdown', True):
|
|
|
- md_path = self._save_markdown(results, doc_output_dir, doc_name)
|
|
|
+ md_path = self._save_markdown(results, middle_json, doc_output_dir, doc_name)
|
|
|
output_paths['markdown'] = str(md_path)
|
|
|
|
|
|
- # 5. 保存表格 HTML
|
|
|
+ # 6. 保存表格 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 模式:保存可视化图片
|
|
|
+ # 7. Debug 模式:保存可视化图片
|
|
|
if output_config.get('save_layout_image', False):
|
|
|
layout_paths = self._save_layout_images(
|
|
|
results, doc_output_dir, doc_name,
|
|
|
@@ -136,20 +163,14 @@ class OutputFormatterV2:
|
|
|
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)
|
|
|
+ # ==================== MinerU middle.json 格式 ====================
|
|
|
|
|
|
def _convert_to_middle_json(self, results: Dict[str, Any]) -> Dict[str, Any]:
|
|
|
- """转换为 MinerU 标准 middle.json 格式"""
|
|
|
+ """
|
|
|
+ 转换为 MinerU 标准 middle.json 格式
|
|
|
+
|
|
|
+ 用于 vlm_union_make 生成 Markdown
|
|
|
+ """
|
|
|
middle_json = {
|
|
|
"pdf_info": [],
|
|
|
"_backend": "vlm",
|
|
|
@@ -167,10 +188,10 @@ class OutputFormatterV2:
|
|
|
}
|
|
|
|
|
|
for element in page.get('elements', []):
|
|
|
- block = self._element_to_block(element, page_info['page_size'])
|
|
|
+ block = self._element_to_middle_block(element)
|
|
|
if block:
|
|
|
elem_type = element.get('type', '')
|
|
|
- if elem_type in ['header', 'footer', 'page_number', 'aside_text']:
|
|
|
+ if elem_type in ['header', 'footer', 'page_number', 'aside_text', 'abandon', 'discarded']:
|
|
|
page_info['discarded_blocks'].append(block)
|
|
|
else:
|
|
|
page_info['para_blocks'].append(block)
|
|
|
@@ -179,12 +200,14 @@ class OutputFormatterV2:
|
|
|
|
|
|
return middle_json
|
|
|
|
|
|
- def _element_to_block(
|
|
|
- self,
|
|
|
- element: Dict[str, Any],
|
|
|
- page_size: List[int]
|
|
|
- ) -> Dict[str, Any]:
|
|
|
- """将元素转换为 MinerU block 格式"""
|
|
|
+ def _element_to_middle_block(self, element: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
|
+ """
|
|
|
+ 将元素转换为 MinerU middle.json block 格式
|
|
|
+
|
|
|
+ MinerU 期望的嵌套结构:
|
|
|
+ - image 类型: { type: "image", blocks: [{ type: "image_body", lines: [...] }] }
|
|
|
+ - table 类型: { type: "table", blocks: [{ type: "table_body", lines: [...] }] }
|
|
|
+ """
|
|
|
elem_type = element.get('type', '')
|
|
|
bbox = element.get('bbox', [0, 0, 0, 0])
|
|
|
content = element.get('content', {})
|
|
|
@@ -209,11 +232,12 @@ class OutputFormatterV2:
|
|
|
}]
|
|
|
}]
|
|
|
|
|
|
- # 表格类型
|
|
|
+ # 表格类型 - 嵌套结构
|
|
|
elif elem_type in ['table', 'table_body']:
|
|
|
table_html = content.get('html', '')
|
|
|
cells = content.get('cells', [])
|
|
|
|
|
|
+ block['type'] = 'table'
|
|
|
block['blocks'] = [{
|
|
|
'type': 'table_body',
|
|
|
'bbox': bbox,
|
|
|
@@ -224,7 +248,26 @@ class OutputFormatterV2:
|
|
|
'bbox': bbox,
|
|
|
'type': 'table',
|
|
|
'html': table_html,
|
|
|
- 'cells': cells # 增强:包含单元格坐标
|
|
|
+ 'cells': cells
|
|
|
+ }]
|
|
|
+ }]
|
|
|
+ }]
|
|
|
+
|
|
|
+ # 图片类型 - 嵌套结构
|
|
|
+ # 注意:MinerU vlm_union_make 期望字段名是 'image_path',不是 'img_path'
|
|
|
+ elif elem_type in ['image', 'image_body', 'figure']:
|
|
|
+ block['type'] = 'image'
|
|
|
+ block['blocks'] = [{
|
|
|
+ 'type': 'image_body',
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'angle': element.get('angle', 0),
|
|
|
+ 'lines': [{
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'spans': [{
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'type': 'image',
|
|
|
+ 'image_path': content.get('image_path', ''), # MinerU 期望 'image_path'
|
|
|
+ 'description': content.get('description', '')
|
|
|
}]
|
|
|
}]
|
|
|
}]
|
|
|
@@ -241,17 +284,273 @@ class OutputFormatterV2:
|
|
|
}]
|
|
|
}]
|
|
|
|
|
|
+ # 表格/图片附属文本
|
|
|
+ elif elem_type in ['table_caption', 'table_footnote', 'image_caption', 'image_footnote']:
|
|
|
+ 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 ['abandon', 'discarded']:
|
|
|
+ block['type'] = 'abandon'
|
|
|
+
|
|
|
return block
|
|
|
|
|
|
+ # ==================== mineru_vllm_results_cell_bbox 格式 ====================
|
|
|
+
|
|
|
+ def _save_page_jsons(
|
|
|
+ self,
|
|
|
+ results: Dict[str, Any],
|
|
|
+ output_dir: Path,
|
|
|
+ doc_name: str
|
|
|
+ ) -> List[str]:
|
|
|
+ """
|
|
|
+ 保存每页独立的 JSON(mineru_vllm_results_cell_bbox 格式)
|
|
|
+
|
|
|
+ 格式示例:
|
|
|
+ [
|
|
|
+ {
|
|
|
+ "type": "table",
|
|
|
+ "img_path": "images/xxx.jpg",
|
|
|
+ "table_caption": [],
|
|
|
+ "table_footnote": [],
|
|
|
+ "table_body": "<table>...</table>",
|
|
|
+ "bbox": [x1, y1, x2, y2],
|
|
|
+ "page_idx": 0,
|
|
|
+ "table_cells": [
|
|
|
+ {
|
|
|
+ "type": "table_cell",
|
|
|
+ "text": "单元格内容",
|
|
|
+ "matched_text": "OCR匹配文本",
|
|
|
+ "bbox": [x1, y1, x2, y2],
|
|
|
+ "row": 1,
|
|
|
+ "col": 1,
|
|
|
+ "score": 100.0,
|
|
|
+ "paddle_bbox_indices": [0, 1]
|
|
|
+ }
|
|
|
+ ],
|
|
|
+ "image_rotation_angle": 0,
|
|
|
+ "skew_angle": 0.0
|
|
|
+ }
|
|
|
+ ]
|
|
|
+ """
|
|
|
+ saved_paths = []
|
|
|
+
|
|
|
+ for page in results.get('pages', []):
|
|
|
+ page_idx = page.get('page_idx', 0)
|
|
|
+ page_name = f"{doc_name}_page_{page_idx + 1:03d}"
|
|
|
+
|
|
|
+ # 转换为 mineru_vllm_results_cell_bbox 格式
|
|
|
+ page_elements = []
|
|
|
+ for element in page.get('elements', []):
|
|
|
+ converted = self._element_to_cell_bbox_format(element, page_idx)
|
|
|
+ if converted:
|
|
|
+ page_elements.append(converted)
|
|
|
+
|
|
|
+ # 保存 JSON
|
|
|
+ json_path = output_dir / f"{page_name}.json"
|
|
|
+ with open(json_path, 'w', encoding='utf-8') as f:
|
|
|
+ json.dump(page_elements, f, ensure_ascii=False, indent=2)
|
|
|
+
|
|
|
+ saved_paths.append(str(json_path))
|
|
|
+ logger.debug(f"📄 Page JSON saved: {json_path}")
|
|
|
+
|
|
|
+ if saved_paths:
|
|
|
+ logger.info(f"📄 {len(saved_paths)} page JSONs saved")
|
|
|
+
|
|
|
+ return saved_paths
|
|
|
+
|
|
|
+ def _element_to_cell_bbox_format(
|
|
|
+ self,
|
|
|
+ element: Dict[str, Any],
|
|
|
+ page_idx: int
|
|
|
+ ) -> Optional[Dict[str, Any]]:
|
|
|
+ """
|
|
|
+ 将元素转换为 mineru_vllm_results_cell_bbox 格式
|
|
|
+ """
|
|
|
+ elem_type = element.get('type', '')
|
|
|
+ bbox = element.get('bbox', [0, 0, 0, 0])
|
|
|
+ content = element.get('content', {})
|
|
|
+
|
|
|
+ # 确保 bbox 是整数列表
|
|
|
+ bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
|
|
|
+
|
|
|
+ result = {
|
|
|
+ 'bbox': bbox,
|
|
|
+ 'page_idx': page_idx
|
|
|
+ }
|
|
|
+
|
|
|
+ # 文本类型
|
|
|
+ if elem_type in ['text', 'title', 'ref_text', 'ocr_text']:
|
|
|
+ text = content.get('text', '') if isinstance(content, dict) else str(content)
|
|
|
+ result['type'] = 'text' if elem_type != 'title' else 'title'
|
|
|
+ result['text'] = text
|
|
|
+ if elem_type == 'title':
|
|
|
+ result['text_level'] = element.get('level', 1)
|
|
|
+
|
|
|
+ # 表格类型
|
|
|
+ elif elem_type in ['table', 'table_body']:
|
|
|
+ result['type'] = 'table'
|
|
|
+ result['img_path'] = content.get('table_image_path', '')
|
|
|
+ result['table_caption'] = self._ensure_list(content.get('table_caption', []))
|
|
|
+ result['table_footnote'] = self._ensure_list(content.get('table_footnote', []))
|
|
|
+ result['table_body'] = content.get('html', '')
|
|
|
+
|
|
|
+ # 关键:table_cells 数组
|
|
|
+ cells = content.get('cells', [])
|
|
|
+ if cells:
|
|
|
+ result['table_cells'] = self._format_table_cells(cells)
|
|
|
+
|
|
|
+ # 旋转和倾斜信息
|
|
|
+ if 'table_angle' in content:
|
|
|
+ result['image_rotation_angle'] = float(content['table_angle'])
|
|
|
+ if 'skew_angle' in content:
|
|
|
+ result['skew_angle'] = float(content['skew_angle'])
|
|
|
+
|
|
|
+ # 图片类型
|
|
|
+ elif elem_type in ['image', 'image_body', 'figure']:
|
|
|
+ result['type'] = 'image'
|
|
|
+ # page JSON 需要完整的相对路径
|
|
|
+ image_filename = content.get('image_path', '')
|
|
|
+ result['img_path'] = f"images/{image_filename}" if image_filename else ''
|
|
|
+ result['image_caption'] = self._ensure_list(content.get('caption', []))
|
|
|
+ result['image_footnote'] = self._ensure_list(content.get('footnote', []))
|
|
|
+
|
|
|
+ # 公式类型
|
|
|
+ elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
|
|
|
+ result['type'] = 'equation'
|
|
|
+ result['text'] = content.get('latex', '') if isinstance(content, dict) else ''
|
|
|
+ result['text_format'] = 'latex'
|
|
|
+
|
|
|
+ # 列表类型
|
|
|
+ elif elem_type == 'list':
|
|
|
+ result['type'] = 'list'
|
|
|
+ result['sub_type'] = 'text'
|
|
|
+ result['list_items'] = content.get('list_items', []) if isinstance(content, dict) else []
|
|
|
+
|
|
|
+ # 页眉页脚
|
|
|
+ elif elem_type in ['header', 'footer']:
|
|
|
+ result['type'] = elem_type
|
|
|
+ result['text'] = content.get('text', '') if isinstance(content, dict) else str(content)
|
|
|
+
|
|
|
+ else:
|
|
|
+ return None
|
|
|
+
|
|
|
+ return result
|
|
|
+
|
|
|
+ def _format_table_cells(self, cells: List[Dict]) -> List[Dict[str, Any]]:
|
|
|
+ """
|
|
|
+ 格式化表格单元格为 mineru_vllm_results_cell_bbox 格式
|
|
|
+
|
|
|
+ 输出格式:
|
|
|
+ {
|
|
|
+ "type": "table_cell",
|
|
|
+ "text": "单元格内容",
|
|
|
+ "matched_text": "OCR匹配文本",
|
|
|
+ "bbox": [x1, y1, x2, y2],
|
|
|
+ "row": 1,
|
|
|
+ "col": 1,
|
|
|
+ "score": 100.0,
|
|
|
+ "paddle_bbox_indices": [0, 1]
|
|
|
+ }
|
|
|
+ """
|
|
|
+ formatted_cells = []
|
|
|
+
|
|
|
+ for cell in cells:
|
|
|
+ formatted_cell = {
|
|
|
+ 'type': 'table_cell',
|
|
|
+ 'text': cell.get('text', ''),
|
|
|
+ 'matched_text': cell.get('matched_text', cell.get('text', '')),
|
|
|
+ 'bbox': [float(x) for x in cell.get('bbox', [0, 0, 0, 0])[:4]],
|
|
|
+ 'row': cell.get('row', 0),
|
|
|
+ 'col': cell.get('col', 0),
|
|
|
+ 'score': float(cell.get('score', 100.0)),
|
|
|
+ 'paddle_bbox_indices': cell.get('paddle_bbox_indices',
|
|
|
+ cell.get('paddle_indices', []))
|
|
|
+ }
|
|
|
+ formatted_cells.append(formatted_cell)
|
|
|
+
|
|
|
+ return formatted_cells
|
|
|
+
|
|
|
+ def _ensure_list(self, value) -> List:
|
|
|
+ """确保值是列表"""
|
|
|
+ if value is None:
|
|
|
+ return []
|
|
|
+ if isinstance(value, str):
|
|
|
+ return [value] if value else []
|
|
|
+ if isinstance(value, list):
|
|
|
+ return value
|
|
|
+ return [str(value)]
|
|
|
+
|
|
|
+ # ==================== Markdown 生成 ====================
|
|
|
+
|
|
|
def _save_markdown(
|
|
|
self,
|
|
|
results: Dict[str, Any],
|
|
|
+ middle_json: Dict[str, Any],
|
|
|
output_dir: Path,
|
|
|
doc_name: str
|
|
|
) -> Path:
|
|
|
- """保存 Markdown 文件"""
|
|
|
+ """
|
|
|
+ 保存 Markdown 文件
|
|
|
+
|
|
|
+ 优先使用 MinerU union_make,降级使用自定义实现
|
|
|
+ """
|
|
|
md_path = output_dir / f"{doc_name}.md"
|
|
|
|
|
|
+ if MINERU_AVAILABLE:
|
|
|
+ try:
|
|
|
+ # image_path 只保存文件名,vlm_union_make 会添加此前缀
|
|
|
+ img_bucket_path = "images"
|
|
|
+ markdown_content = vlm_union_make(
|
|
|
+ middle_json['pdf_info'],
|
|
|
+ MakeMode.MM_MD,
|
|
|
+ img_bucket_path
|
|
|
+ )
|
|
|
+
|
|
|
+ if markdown_content:
|
|
|
+ if isinstance(markdown_content, list):
|
|
|
+ markdown_content = '\n\n'.join(markdown_content)
|
|
|
+
|
|
|
+ header = self._generate_markdown_header(results)
|
|
|
+ markdown_content = header + str(markdown_content)
|
|
|
+
|
|
|
+ with open(md_path, 'w', encoding='utf-8') as f:
|
|
|
+ f.write(markdown_content)
|
|
|
+
|
|
|
+ logger.info(f"📝 Markdown saved (MinerU format): {md_path}")
|
|
|
+ return md_path
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ logger.warning(f"MinerU union_make failed: {e}, falling back to custom implementation")
|
|
|
+
|
|
|
+ # 降级方案
|
|
|
+ markdown_content = self._generate_markdown_fallback(results)
|
|
|
+ with open(md_path, 'w', encoding='utf-8') as f:
|
|
|
+ f.write(markdown_content)
|
|
|
+
|
|
|
+ logger.info(f"📝 Markdown saved (fallback): {md_path}")
|
|
|
+ return md_path
|
|
|
+
|
|
|
+ def _generate_markdown_header(self, results: Dict[str, Any]) -> str:
|
|
|
+ """生成 Markdown 文件头"""
|
|
|
+ return f"""---
|
|
|
+scene: {results.get('scene', 'unknown')}
|
|
|
+document: {results.get('document_path', '')}
|
|
|
+pages: {len(results.get('pages', []))}
|
|
|
+---
|
|
|
+
|
|
|
+"""
|
|
|
+
|
|
|
+ def _generate_markdown_fallback(self, results: Dict[str, Any]) -> str:
|
|
|
+ """降级方案:自定义 Markdown 生成"""
|
|
|
md_lines = [
|
|
|
f"---",
|
|
|
f"scene: {results.get('scene', 'unknown')}",
|
|
|
@@ -263,7 +562,6 @@ class OutputFormatterV2:
|
|
|
|
|
|
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', '')
|
|
|
@@ -271,83 +569,92 @@ class OutputFormatterV2:
|
|
|
bbox = element.get('bbox', [])
|
|
|
|
|
|
# 添加 bbox 注释
|
|
|
- bbox_comment = f"<!-- bbox: {bbox} -->"
|
|
|
+ if bbox:
|
|
|
+ md_lines.append(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']:
|
|
|
+ # 表格标题
|
|
|
+ for caption in self._ensure_list(content.get('table_caption', [])):
|
|
|
+ md_lines.append(f"**{caption}**")
|
|
|
+
|
|
|
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(f"\n{html}\n")
|
|
|
+ md_lines.append("")
|
|
|
+
|
|
|
+ elif elem_type in ['image', 'image_body', 'figure']:
|
|
|
+ img_filename = content.get('image_path', '')
|
|
|
+ if img_filename:
|
|
|
+ md_lines.append(f"")
|
|
|
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
|
|
|
+ return '\n'.join(md_lines)
|
|
|
|
|
|
- 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)
|
|
|
+ # ==================== 图片保存 ====================
|
|
|
+
|
|
|
+ def _save_image_elements(
|
|
|
+ self,
|
|
|
+ results: Dict[str, Any],
|
|
|
+ images_dir: Path,
|
|
|
+ doc_name: str
|
|
|
+ ) -> List[str]:
|
|
|
+ """保存图片元素"""
|
|
|
+ saved_paths = []
|
|
|
+ image_count = 0
|
|
|
+
|
|
|
+ for page in results.get('pages', []):
|
|
|
+ page_idx = page.get('page_idx', 0)
|
|
|
|
|
|
- except Exception as e:
|
|
|
- logger.warning(f"HTML to Markdown conversion failed: {e}")
|
|
|
- return html
|
|
|
+ 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)
|
|
|
+
|
|
|
+ # 更新路径(只保存文件名,不含 images/ 前缀)
|
|
|
+ # vlm_union_make 会自动添加 img_bucket_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
|
|
|
+
|
|
|
+ # ==================== 表格 HTML ====================
|
|
|
|
|
|
def _save_table_htmls(
|
|
|
self,
|
|
|
@@ -355,7 +662,7 @@ class OutputFormatterV2:
|
|
|
output_dir: Path,
|
|
|
doc_name: str
|
|
|
) -> Path:
|
|
|
- """保存表格 HTML 文件(带坐标信息)"""
|
|
|
+ """保存表格 HTML 文件"""
|
|
|
tables_dir = output_dir / 'tables'
|
|
|
tables_dir.mkdir(exist_ok=True)
|
|
|
|
|
|
@@ -372,7 +679,6 @@ class OutputFormatterV2:
|
|
|
cells = content.get('cells', [])
|
|
|
|
|
|
if html:
|
|
|
- # 生成带样式的 HTML
|
|
|
full_html = self._generate_table_html_with_styles(
|
|
|
html, cells, doc_name, page_idx, table_count
|
|
|
)
|
|
|
@@ -394,7 +700,7 @@ class OutputFormatterV2:
|
|
|
page_idx: int,
|
|
|
table_idx: int
|
|
|
) -> str:
|
|
|
- """生成带样式和坐标信息的完整 HTML"""
|
|
|
+ """生成带样式的完整 HTML"""
|
|
|
cells_json = json.dumps(cells, ensure_ascii=False, indent=2) if cells else "[]"
|
|
|
|
|
|
return f"""<!DOCTYPE html>
|
|
|
@@ -499,6 +805,8 @@ class OutputFormatterV2:
|
|
|
</body>
|
|
|
</html>"""
|
|
|
|
|
|
+ # ==================== Debug 可视化 ====================
|
|
|
+
|
|
|
def _save_layout_images(
|
|
|
self,
|
|
|
results: Dict[str, Any],
|
|
|
@@ -512,12 +820,14 @@ class OutputFormatterV2:
|
|
|
|
|
|
for page in results.get('pages', []):
|
|
|
page_idx = page.get('page_idx', 0)
|
|
|
- processed_image = page.get('processed_image')
|
|
|
+ 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
|
|
|
|
|
|
- # 转换为 PIL Image
|
|
|
if isinstance(processed_image, np.ndarray):
|
|
|
image = Image.fromarray(processed_image).convert('RGB')
|
|
|
elif isinstance(processed_image, Image.Image):
|
|
|
@@ -525,13 +835,9 @@ class OutputFormatterV2:
|
|
|
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])
|
|
|
@@ -542,31 +848,30 @@ class OutputFormatterV2:
|
|
|
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))
|
|
|
@@ -580,17 +885,19 @@ class OutputFormatterV2:
|
|
|
output_dir: Path,
|
|
|
doc_name: str
|
|
|
) -> List[str]:
|
|
|
- """保存 OCR 可视化图片(显示文本框和单元格坐标)"""
|
|
|
+ """保存 OCR 可视化图片"""
|
|
|
ocr_paths = []
|
|
|
|
|
|
for page in results.get('pages', []):
|
|
|
page_idx = page.get('page_idx', 0)
|
|
|
- processed_image = page.get('processed_image')
|
|
|
+ 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
|
|
|
|
|
|
- # 转换为 PIL Image
|
|
|
if isinstance(processed_image, np.ndarray):
|
|
|
image = Image.fromarray(processed_image).convert('RGB')
|
|
|
elif isinstance(processed_image, Image.Image):
|
|
|
@@ -601,18 +908,17 @@ class OutputFormatterV2:
|
|
|
draw = ImageDraw.Draw(image)
|
|
|
font = self._get_font(10)
|
|
|
|
|
|
- # 遍历所有元素
|
|
|
for element in page.get('elements', []):
|
|
|
content = element.get('content', {})
|
|
|
|
|
|
- # 绘制 OCR 文本框
|
|
|
+ # 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', [])
|
|
|
@@ -620,19 +926,17 @@ class OutputFormatterV2:
|
|
|
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 框
|
|
|
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))
|
|
|
@@ -648,12 +952,10 @@ class OutputFormatterV2:
|
|
|
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]) # 闭合多边形
|
|
|
+ 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)
|
|
|
@@ -675,6 +977,38 @@ class OutputFormatterV2:
|
|
|
return ImageFont.load_default()
|
|
|
|
|
|
|
|
|
+# ==================== 便捷函数 ====================
|
|
|
+
|
|
|
+def save_mineru_format(
|
|
|
+ results: Dict[str, Any],
|
|
|
+ output_dir: str,
|
|
|
+ output_config: Dict[str, Any] = None
|
|
|
+) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 便捷函数:保存为 MinerU 格式
|
|
|
+
|
|
|
+ Args:
|
|
|
+ results: pipeline 处理结果
|
|
|
+ output_dir: 输出目录
|
|
|
+ output_config: 输出配置
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 输出文件路径字典
|
|
|
+ """
|
|
|
+ if output_config is None:
|
|
|
+ output_config = {
|
|
|
+ 'save_json': True,
|
|
|
+ 'save_page_json': True,
|
|
|
+ 'save_markdown': True,
|
|
|
+ 'save_html': True,
|
|
|
+ 'save_layout_image': False,
|
|
|
+ 'save_ocr_image': False,
|
|
|
+ }
|
|
|
+
|
|
|
+ formatter = OutputFormatterV2(output_dir)
|
|
|
+ return formatter.save_results(results, output_config)
|
|
|
+
|
|
|
+
|
|
|
if __name__ == "__main__":
|
|
|
# 测试代码
|
|
|
sample_results = {
|
|
|
@@ -707,11 +1041,12 @@ if __name__ == "__main__":
|
|
|
]
|
|
|
}
|
|
|
|
|
|
- formatter = OutputFormatterV2("./test_output_v2")
|
|
|
- output_files = formatter.save_results(
|
|
|
+ output_files = save_mineru_format(
|
|
|
sample_results,
|
|
|
+ "./test_output_v2",
|
|
|
{
|
|
|
"save_json": True,
|
|
|
+ "save_page_json": True,
|
|
|
"save_markdown": True,
|
|
|
"save_html": True,
|
|
|
"save_layout_image": False,
|