"""
统一输出格式化器
严格遵循 MinerU mineru_vllm_results_cell_bbox 格式
支持:
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, Tuple
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 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 = {
'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, Any]:
"""
保存处理结果
Args:
results: 处理结果
output_config: 输出配置
Returns:
输出文件路径字典
"""
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)
# 创建 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)
# 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:
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}")
# 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
# 5. 保存 Markdown
if output_config.get('save_markdown', True):
md_path = self._save_markdown(results, middle_json, doc_output_dir, doc_name)
output_paths['markdown'] = str(md_path)
# 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)
# 7. 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
# ==================== MinerU middle.json 格式 ====================
def _convert_to_middle_json(self, results: Dict[str, Any]) -> Dict[str, Any]:
"""
转换为 MinerU 标准 middle.json 格式
用于 vlm_union_make 生成 Markdown
"""
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_middle_block(element)
if block:
elem_type = element.get('type', '')
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)
middle_json['pdf_info'].append(page_info)
return middle_json
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', {})
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['type'] = 'table'
block['blocks'] = [{
'type': 'table_body',
'bbox': bbox,
'angle': 0,
'lines': [{
'bbox': bbox,
'spans': [{
'bbox': bbox,
'type': 'table',
'html': table_html,
'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', '')
}]
}]
}]
# 公式类型
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
}]
}]
# 表格/图片附属文本
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": "
",
"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 文件
优先使用 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')}",
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)
for element in page.get('elements', []):
elem_type = element.get('type', '')
content = element.get('content', {})
bbox = element.get('bbox', [])
# 添加 bbox 注释
if bbox:
md_lines.append(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("")
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("")
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:
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("")
return '\n'.join(md_lines)
# ==================== 图片保存 ====================
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)
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,
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:
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}
{table_html}
📍 单元格坐标数据 (JSON)
{cells_json}
"""
# ==================== Debug 可视化 ====================
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('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 = 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('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 = 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
):
"""绘制多边形或矩形"""
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)
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()
# ==================== 便捷函数 ====================
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 = {
"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}
]
}
}
]
}
]
}
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,
"save_ocr_image": False
}
)
print("Generated files:", output_files)