| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058 |
- """
- 统一输出格式化器
- 严格遵循 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": "<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 文件
-
- 优先使用 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"<!-- 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("")
-
- 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"""<!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>"""
-
- # ==================== 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": "<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}
- ]
- }
- }
- ]
- }
- ]
- }
-
- 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)
|