init: RAG 知识库服务初始提交
- 后端 API(Flask + Gunicorn) - RAG 引擎(混合检索 + 云端 Reranker + 引用溯源) - 文档解析(MinerU + 多格式支持) - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
365
parsers/__init__.py
Normal file
365
parsers/__init__.py
Normal file
@@ -0,0 +1,365 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
文档解析器模块 (v5 - MinerU 统一版)
|
||||
|
||||
统一入口:parse_document(filepath) -> List[UnifiedChunk]
|
||||
|
||||
格式支持:
|
||||
- PDF/DOCX/PPTX/图片 → parse_with_mineru_persistent()
|
||||
- XLSX/XLS → parse_excel() (Pandas 专属管道)
|
||||
- TXT → parse_txt()
|
||||
|
||||
MinerU 3.0+ 优势:
|
||||
- PDF: 表格识别率 95%+,支持 109 种语言 OCR
|
||||
- DOCX: 原生解析,速度提升数十倍,无幻觉
|
||||
- 图片自动提取,路径存入 UnifiedChunk
|
||||
|
||||
依赖:
|
||||
pip install "mineru[all]"
|
||||
pip install pandas openpyxl
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ========== 支持的文件格式 ==========
|
||||
|
||||
SUPPORTED_FORMATS = {
|
||||
# MinerU 支持
|
||||
'.pdf': 'PDF 文档',
|
||||
'.docx': 'Word 文档',
|
||||
'.pptx': 'PowerPoint 幻灯片',
|
||||
'.png': 'PNG 图片',
|
||||
'.jpg': 'JPEG 图片',
|
||||
'.jpeg': 'JPEG 图片',
|
||||
'.bmp': 'BMP 图片',
|
||||
'.tiff': 'TIFF 图片',
|
||||
# Pandas 支持
|
||||
'.xlsx': 'Excel 表格',
|
||||
'.xls': 'Excel 表格',
|
||||
# 文本
|
||||
'.txt': '文本文件',
|
||||
}
|
||||
|
||||
# ========== 模块可用性检测 ==========
|
||||
|
||||
MINERU_AVAILABLE = False
|
||||
PANDAS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from parsers.mineru_parser import (
|
||||
parse_with_mineru_persistent,
|
||||
parse_with_mineru,
|
||||
MinerUChunk,
|
||||
convert_to_rag_format as mineru_to_rag_format,
|
||||
)
|
||||
MINERU_AVAILABLE = True
|
||||
except ImportError as e:
|
||||
logger.warning(f"MinerU 不可用: {e}")
|
||||
|
||||
try:
|
||||
from parsers.excel_parser import (
|
||||
parse_excel,
|
||||
get_table_meta,
|
||||
convert_to_rag_format as excel_to_rag_format,
|
||||
UnifiedChunk as ExcelChunk,
|
||||
)
|
||||
PANDAS_AVAILABLE = True
|
||||
except ImportError as e:
|
||||
logger.warning(f"Excel 解析器不可用: {e}")
|
||||
|
||||
try:
|
||||
from parsers.txt_parser import extract_text_from_txt
|
||||
TXT_AVAILABLE = True
|
||||
except ImportError:
|
||||
TXT_AVAILABLE = False
|
||||
|
||||
|
||||
# ========== 统一 Schema ==========
|
||||
|
||||
@dataclass
|
||||
class UnifiedChunk:
|
||||
"""
|
||||
统一内部 Schema - 所有解析器输出此格式
|
||||
|
||||
与 MinerUChunk 完全兼容,方便下游处理。
|
||||
"""
|
||||
content: str # 文本内容(Markdown 格式)
|
||||
chunk_type: str # 类型: text, table, image, equation
|
||||
page_start: int = 1 # 起始页码/行号
|
||||
page_end: int = 1 # 结束页码/行号
|
||||
text_level: int = 0 # 标题级别 (0=body, 1=h1, ...)
|
||||
title: str = "" # 标题文本
|
||||
section_path: str = "" # 章节路径
|
||||
source_file: str = "" # 源文件名
|
||||
bbox: Optional[List[float]] = None # 边界框 [x0, y0, x1, y1]
|
||||
table_html: Optional[str] = None # 表格 HTML(表格类型)
|
||||
image_path: Optional[str] = None # 图片路径(图片类型)
|
||||
|
||||
|
||||
class UnsupportedFormatError(Exception):
|
||||
"""不支持的文件格式异常"""
|
||||
pass
|
||||
|
||||
|
||||
# ========== 统一入口函数 ==========
|
||||
|
||||
def parse_document(
|
||||
filepath: str,
|
||||
output_base: str = ".data/mineru_temp",
|
||||
images_output: str = ".data/images",
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
统一文档解析入口(扁平化存储)
|
||||
|
||||
Args:
|
||||
filepath: 文档文件路径
|
||||
output_base: MinerU 临时输出目录
|
||||
images_output: 图片存储目录
|
||||
**kwargs: 格式特定参数
|
||||
|
||||
Returns:
|
||||
{
|
||||
'chunks': List[UnifiedChunk], # 结构化分块
|
||||
'markdown': str, # Markdown 内容
|
||||
'tables': List[str], # 表格列表
|
||||
'images': List[str], # 图片列表
|
||||
'source_file': str, # 源文件名
|
||||
'parser_used': str, # 使用的解析器
|
||||
}
|
||||
|
||||
Raises:
|
||||
UnsupportedFormatError: 不支持的文件格式
|
||||
FileNotFoundError: 文件不存在
|
||||
"""
|
||||
filepath = Path(filepath)
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {filepath}")
|
||||
|
||||
ext = filepath.suffix.lower()
|
||||
if ext not in SUPPORTED_FORMATS:
|
||||
raise UnsupportedFormatError(
|
||||
f"不支持的文件格式: {ext}。"
|
||||
f"支持格式: {', '.join(SUPPORTED_FORMATS.keys())}"
|
||||
)
|
||||
|
||||
logger.info(f"解析 {SUPPORTED_FORMATS.get(ext, '文档')}: {filepath.name}")
|
||||
|
||||
# 根据扩展名选择解析器
|
||||
if ext in ('.pdf', '.docx', '.pptx', '.png', '.jpg', '.jpeg', '.bmp', '.tiff'):
|
||||
return _parse_with_mineru(filepath, output_base, images_output, **kwargs)
|
||||
elif ext in ('.xlsx', '.xls'):
|
||||
return _parse_with_pandas(filepath, **kwargs)
|
||||
elif ext == '.txt':
|
||||
return _parse_txt(filepath, **kwargs)
|
||||
else:
|
||||
raise UnsupportedFormatError(f"不支持的文件格式: {ext}")
|
||||
|
||||
|
||||
def _parse_with_mineru(
|
||||
filepath: Path,
|
||||
output_base: str,
|
||||
images_output: str,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""使用 MinerU 解析文档"""
|
||||
if not MINERU_AVAILABLE:
|
||||
raise RuntimeError("MinerU 不可用,请运行: pip install \"mineru[all]\"")
|
||||
|
||||
result = parse_with_mineru_persistent(
|
||||
str(filepath),
|
||||
output_base=output_base,
|
||||
images_output=images_output,
|
||||
cleanup_after_image_move=kwargs.get('cleanup_after_image_move', False)
|
||||
)
|
||||
|
||||
# 转换 chunks 为 UnifiedChunk 格式(已是 MinerUChunk,兼容)
|
||||
chunks = result.get('chunks', [])
|
||||
|
||||
return {
|
||||
'chunks': chunks,
|
||||
'markdown': result.get('markdown', ''),
|
||||
'tables': result.get('tables', []),
|
||||
'images': result.get('images', []),
|
||||
'source_file': filepath.name,
|
||||
'parser_used': 'mineru',
|
||||
'file_hash': result.get('file_hash', ''),
|
||||
'output_dir': result.get('output_dir', ''),
|
||||
}
|
||||
|
||||
|
||||
def _parse_with_pandas(filepath: Path, **kwargs) -> Dict[str, Any]:
|
||||
"""使用 Pandas 解析 Excel"""
|
||||
if not PANDAS_AVAILABLE:
|
||||
raise RuntimeError("Excel 解析器不可用,请运行: pip install pandas openpyxl")
|
||||
|
||||
result = parse_excel(
|
||||
str(filepath),
|
||||
max_rows_per_chunk=kwargs.get('max_rows_per_chunk', 200)
|
||||
)
|
||||
|
||||
# 转换 chunks 为 UnifiedChunk 格式(已是 UnifiedChunk)
|
||||
chunks = result.get('chunks', [])
|
||||
|
||||
# 构建 Markdown
|
||||
markdown_parts = []
|
||||
for chunk in chunks:
|
||||
markdown_parts.append(f"## {chunk.title}\n\n{chunk.content}\n")
|
||||
|
||||
return {
|
||||
'chunks': chunks,
|
||||
'markdown': "\n".join(markdown_parts),
|
||||
'tables': [chunk.content for chunk in chunks],
|
||||
'images': [],
|
||||
'source_file': filepath.name,
|
||||
'parser_used': 'pandas',
|
||||
'sheets': result.get('sheets', []),
|
||||
'total_rows': result.get('total_rows', 0),
|
||||
}
|
||||
|
||||
|
||||
def _parse_txt(filepath: Path, **kwargs) -> Dict[str, Any]:
|
||||
"""解析纯文本文件"""
|
||||
# 直接读取文件内容
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
|
||||
# 简单分块
|
||||
chunk_size = kwargs.get('chunk_size', 1000)
|
||||
chunks = []
|
||||
|
||||
for i in range(0, len(content), chunk_size):
|
||||
chunk_content = content[i:i+chunk_size]
|
||||
chunk = UnifiedChunk(
|
||||
content=chunk_content,
|
||||
chunk_type="text",
|
||||
page_start=i // chunk_size + 1,
|
||||
page_end=i // chunk_size + 1,
|
||||
source_file=filepath.name
|
||||
)
|
||||
chunks.append(chunk)
|
||||
|
||||
return {
|
||||
'chunks': chunks,
|
||||
'markdown': content,
|
||||
'tables': [],
|
||||
'images': [],
|
||||
'source_file': filepath.name,
|
||||
'parser_used': 'txt',
|
||||
}
|
||||
|
||||
|
||||
# ========== RAG 格式转换 ==========
|
||||
|
||||
def convert_to_rag_format(result: Dict[str, Any]) -> List[Dict]:
|
||||
"""
|
||||
将解析结果转换为 RAG 入库格式
|
||||
|
||||
Args:
|
||||
result: parse_document() 返回结果
|
||||
|
||||
Returns:
|
||||
[{'text': ..., 'page': ..., 'has_table': ..., ...}, ...]
|
||||
"""
|
||||
parser_used = result.get('parser_used', 'unknown')
|
||||
chunks = result.get('chunks', [])
|
||||
|
||||
if parser_used == 'mineru':
|
||||
# MinerU chunks 已有专用转换函数
|
||||
from parsers.mineru_parser import convert_to_rag_format as mineru_convert
|
||||
return mineru_convert(result, result.get('source_file', ''))
|
||||
|
||||
elif parser_used == 'pandas':
|
||||
# Excel chunks
|
||||
from parsers.excel_parser import convert_to_rag_format as excel_convert
|
||||
return excel_convert(result)
|
||||
|
||||
else:
|
||||
# 通用转换
|
||||
pages_content = []
|
||||
for chunk in chunks:
|
||||
page_info = {
|
||||
'text': chunk.content,
|
||||
'page': chunk.page_start,
|
||||
'page_end': chunk.page_end,
|
||||
'has_table': chunk.chunk_type == 'table',
|
||||
'section': chunk.title,
|
||||
'section_path': chunk.section_path,
|
||||
'level': chunk.text_level,
|
||||
'chunk_type': chunk.chunk_type,
|
||||
'source_file': chunk.source_file,
|
||||
}
|
||||
pages_content.append(page_info)
|
||||
|
||||
return pages_content
|
||||
|
||||
|
||||
# ========== 兼容旧接口 ==========
|
||||
|
||||
def extract_text_from_pdf(filepath, **kwargs):
|
||||
"""兼容旧接口:从 PDF 提取文本"""
|
||||
result = parse_document(filepath, **kwargs)
|
||||
pages_content = convert_to_rag_format(result)
|
||||
images_info = [{'id': img} for img in result.get('images', [])]
|
||||
return pages_content, images_info
|
||||
|
||||
|
||||
def extract_text_from_docx(filepath, **kwargs):
|
||||
"""兼容旧接口:从 Word 提取文本"""
|
||||
result = parse_document(filepath, **kwargs)
|
||||
return convert_to_rag_format(result)
|
||||
|
||||
|
||||
def extract_text_from_xlsx(filepath, **kwargs):
|
||||
"""兼容旧接口:从 Excel 提取文本"""
|
||||
result = parse_document(filepath, **kwargs)
|
||||
return convert_to_rag_format(result)
|
||||
|
||||
|
||||
def extract_text_from_txt(filepath, **kwargs):
|
||||
"""兼容旧接口:从 TXT 提取文本"""
|
||||
# 直接读取文件,避免递归调用 parse_document
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
|
||||
# 简单分块
|
||||
chunk_size = kwargs.get('chunk_size', 1000)
|
||||
chunks = []
|
||||
for i in range(0, len(content), chunk_size):
|
||||
chunks.append({
|
||||
"content": content[i:i+chunk_size],
|
||||
"chunk_type": "text",
|
||||
"page": i // chunk_size + 1
|
||||
})
|
||||
|
||||
return {
|
||||
"chunks": chunks,
|
||||
"markdown": content,
|
||||
"tables": [],
|
||||
"images": []
|
||||
}
|
||||
|
||||
|
||||
# ========== 模块导出 ==========
|
||||
|
||||
__all__ = [
|
||||
# 统一入口
|
||||
'parse_document',
|
||||
'convert_to_rag_format',
|
||||
'UnifiedChunk',
|
||||
'UnsupportedFormatError',
|
||||
'SUPPORTED_FORMATS',
|
||||
# 兼容旧接口
|
||||
'extract_text_from_pdf',
|
||||
'extract_text_from_docx',
|
||||
'extract_text_from_xlsx',
|
||||
'extract_text_from_txt',
|
||||
# 可用性标志
|
||||
'MINERU_AVAILABLE',
|
||||
'PANDAS_AVAILABLE',
|
||||
]
|
||||
474
parsers/excel_parser.py
Normal file
474
parsers/excel_parser.py
Normal file
@@ -0,0 +1,474 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Excel 解析模块(Pandas 管道)
|
||||
|
||||
MinerU 不支持 XLSX 格式,因此使用 Pandas 专属管道处理。
|
||||
|
||||
策略:表级摘要(Chroma)+ 完整 Markdown(DocStore)
|
||||
- 每个 sheet 生成 Markdown 表格
|
||||
- 大表(>200行)按行切片,每片保留表头
|
||||
- 每片包装为 UnifiedChunk(type='table')
|
||||
|
||||
检索链路:
|
||||
用户提问 → 命中摘要 → 拿 doc_id → 掏 Markdown → 喂 LLM
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional, Any
|
||||
from dataclasses import dataclass, field
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 大表切片阈值
|
||||
MAX_ROWS_PER_CHUNK = 200
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnifiedChunk:
|
||||
"""统一内部 Schema - 与 MinerUChunk 兼容"""
|
||||
content: str # 文本内容(Markdown 格式)
|
||||
chunk_type: str # 类型: table
|
||||
page_start: int = 1 # 起始行号(Excel 无页码概念)
|
||||
page_end: int = 1 # 结束行号
|
||||
text_level: int = 0 # 标题级别(Excel 无标题层级)
|
||||
title: str = "" # Sheet 名称
|
||||
section_path: str = "" # 章节路径
|
||||
source_file: str = "" # 源文件名
|
||||
bbox: Optional[List[float]] = field(default=None) # 不适用
|
||||
table_html: Optional[str] = field(default=None) # 表格 HTML(可选)
|
||||
image_path: Optional[str] = field(default=None) # 不适用
|
||||
# Excel 专用元数据
|
||||
sheet_name: str = "" # Sheet 名称
|
||||
row_start: int = 0 # 起始行(0-indexed)
|
||||
row_end: int = 0 # 结束行
|
||||
col_count: int = 0 # 列数
|
||||
headers: List[str] = field(default_factory=list) # 表头列表
|
||||
|
||||
|
||||
def parse_excel(
|
||||
filepath: str,
|
||||
max_rows_per_chunk: int = MAX_ROWS_PER_CHUNK
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
解析 Excel 文件,输出 UnifiedChunk 列表
|
||||
|
||||
Args:
|
||||
filepath: Excel 文件路径
|
||||
max_rows_per_chunk: 大表切片阈值,默认 200 行
|
||||
|
||||
Returns:
|
||||
{
|
||||
'chunks': List[UnifiedChunk], # 结构化分块
|
||||
'sheets': List[str], # Sheet 名称列表
|
||||
'total_rows': int, # 总行数
|
||||
'source_file': str # 源文件名
|
||||
}
|
||||
"""
|
||||
filepath = Path(filepath)
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {filepath}")
|
||||
|
||||
logger.info(f"使用 Pandas 解析 Excel: {filepath.name}")
|
||||
|
||||
chunks = []
|
||||
sheet_names = []
|
||||
total_rows = 0
|
||||
|
||||
# 读取所有 sheets
|
||||
try:
|
||||
xls = pd.ExcelFile(filepath)
|
||||
sheet_names = xls.sheet_names
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Excel 文件读取失败: {e}")
|
||||
|
||||
for sheet_name in sheet_names:
|
||||
try:
|
||||
# 先读取原始数据(不指定表头)
|
||||
df_raw = pd.read_excel(filepath, sheet_name=sheet_name, header=None)
|
||||
except Exception as e:
|
||||
logger.warning(f"Sheet '{sheet_name}' 读取失败: {e}")
|
||||
continue
|
||||
|
||||
if df_raw.empty:
|
||||
logger.debug(f"Sheet '{sheet_name}' 为空,跳过")
|
||||
continue
|
||||
|
||||
# 清理数据:填充 NaN
|
||||
df_raw = df_raw.fillna('')
|
||||
|
||||
# 检测表头行(查找包含"部门"、"负责人"等典型表头关键词的行)
|
||||
header_row_idx = _detect_header_row(df_raw)
|
||||
|
||||
# 提取表格标题(表头上方的行)
|
||||
table_title = ""
|
||||
if header_row_idx > 0:
|
||||
# 表头上方的第一行可能是标题
|
||||
first_row = df_raw.iloc[0]
|
||||
first_row_text = ' '.join([str(v) for v in first_row if str(v).strip()])
|
||||
if first_row_text and len(first_row_text) < 50:
|
||||
table_title = first_row_text
|
||||
|
||||
# 重新读取,使用检测到的表头行
|
||||
if header_row_idx is not None and header_row_idx > 0:
|
||||
df = pd.read_excel(filepath, sheet_name=sheet_name, header=header_row_idx)
|
||||
else:
|
||||
df = pd.read_excel(filepath, sheet_name=sheet_name)
|
||||
|
||||
df = df.fillna('')
|
||||
|
||||
row_count = len(df)
|
||||
col_count = len(df.columns)
|
||||
total_rows += row_count
|
||||
|
||||
# 获取表头
|
||||
headers = [str(col) for col in df.columns.tolist()]
|
||||
|
||||
# 过滤掉 Unnamed 列名
|
||||
headers = [h if not h.startswith('Unnamed') else f'列{i+1}' for i, h in enumerate(headers)]
|
||||
|
||||
# 大表切片
|
||||
if row_count > max_rows_per_chunk:
|
||||
logger.info(f"Sheet '{sheet_name}' 有 {row_count} 行,按 {max_rows_per_chunk} 行切片")
|
||||
|
||||
num_chunks = (row_count + max_rows_per_chunk - 1) // max_rows_per_chunk
|
||||
|
||||
for i in range(num_chunks):
|
||||
start_row = i * max_rows_per_chunk
|
||||
end_row = min((i + 1) * max_rows_per_chunk, row_count)
|
||||
|
||||
# 切片数据(保留表头)
|
||||
df_slice = df.iloc[start_row:end_row]
|
||||
|
||||
# 转 Markdown
|
||||
md_table = _df_to_markdown(df_slice, headers)
|
||||
|
||||
# 标注切片信息
|
||||
chunk_title = f"{sheet_name} (第{i+1}/{num_chunks}片,行{start_row+1}-{end_row})"
|
||||
|
||||
chunk = UnifiedChunk(
|
||||
content=md_table,
|
||||
chunk_type="table",
|
||||
page_start=start_row + 1,
|
||||
page_end=end_row,
|
||||
title=chunk_title,
|
||||
section_path=sheet_name,
|
||||
source_file=filepath.name,
|
||||
sheet_name=sheet_name,
|
||||
row_start=start_row,
|
||||
row_end=end_row,
|
||||
col_count=col_count,
|
||||
headers=headers
|
||||
)
|
||||
chunks.append(chunk)
|
||||
else:
|
||||
# 小表直接转换
|
||||
md_table = _df_to_markdown(df, headers)
|
||||
|
||||
# 使用表格标题或 sheet 名称
|
||||
chunk_title = table_title if table_title else sheet_name
|
||||
|
||||
chunk = UnifiedChunk(
|
||||
content=md_table,
|
||||
chunk_type="table",
|
||||
page_start=1,
|
||||
page_end=row_count,
|
||||
title=chunk_title,
|
||||
section_path=sheet_name,
|
||||
source_file=filepath.name,
|
||||
sheet_name=sheet_name,
|
||||
row_start=0,
|
||||
row_end=row_count,
|
||||
col_count=col_count,
|
||||
headers=headers
|
||||
)
|
||||
chunks.append(chunk)
|
||||
|
||||
logger.info(f"Excel 解析完成: {len(chunks)} 个表格块,{total_rows} 行数据")
|
||||
|
||||
return {
|
||||
'chunks': chunks,
|
||||
'sheets': sheet_names,
|
||||
'total_rows': total_rows,
|
||||
'source_file': filepath.name
|
||||
}
|
||||
|
||||
|
||||
def _detect_header_row(df: pd.DataFrame) -> Optional[int]:
|
||||
"""
|
||||
检测表头行位置
|
||||
|
||||
表头特征:
|
||||
1. 包含典型表头关键词(部门、负责人、名称、数量等)
|
||||
2. 不含大量数字(数据行特征)
|
||||
3. 文本较短
|
||||
|
||||
Returns:
|
||||
表头行索引(0-indexed),未找到返回 0
|
||||
"""
|
||||
# 典型表头关键词
|
||||
header_keywords = {
|
||||
'部门', '负责人', '名称', '数量', '人数', '金额', '日期', '地址',
|
||||
'电话', '邮箱', '编号', '类型', '状态', '备注', '描述', '职位',
|
||||
'团队', '职责', '地点', '公司', '分', '总', '规模', '职能'
|
||||
}
|
||||
|
||||
best_row = 0
|
||||
best_score = 0
|
||||
|
||||
for idx in range(min(5, len(df))): # 只检查前 5 行
|
||||
row = df.iloc[idx]
|
||||
score = 0
|
||||
|
||||
for cell in row:
|
||||
cell_str = str(cell).strip()
|
||||
if not cell_str:
|
||||
continue
|
||||
|
||||
# 检查关键词
|
||||
for kw in header_keywords:
|
||||
if kw in cell_str:
|
||||
score += 2
|
||||
|
||||
# 短文本倾向于表头
|
||||
if len(cell_str) < 20:
|
||||
score += 1
|
||||
|
||||
# 数字倾向于数据行
|
||||
try:
|
||||
float(cell_str)
|
||||
score -= 3
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_row = idx
|
||||
|
||||
return best_row
|
||||
|
||||
|
||||
def _df_to_markdown(df: pd.DataFrame, headers: List[str] = None) -> str:
|
||||
"""
|
||||
将 DataFrame 转换为 Markdown 表格格式
|
||||
|
||||
Args:
|
||||
df: DataFrame
|
||||
headers: 表头列表(可选,默认使用 df.columns)
|
||||
|
||||
Returns:
|
||||
Markdown 表格字符串
|
||||
"""
|
||||
if headers is None:
|
||||
headers = [str(col) for col in df.columns.tolist()]
|
||||
|
||||
lines = []
|
||||
|
||||
# 表头行
|
||||
header_line = "| " + " | ".join(headers) + " |"
|
||||
lines.append(header_line)
|
||||
|
||||
# 分隔行
|
||||
separator = "| " + " | ".join(["---"] * len(headers)) + " |"
|
||||
lines.append(separator)
|
||||
|
||||
# 数据行
|
||||
for _, row in df.iterrows():
|
||||
cells = [str(val).replace('\n', ' ').replace('|', '\\|') for val in row]
|
||||
data_line = "| " + " | ".join(cells) + " |"
|
||||
lines.append(data_line)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def get_table_meta(filepath: str, sheet_name: str = None) -> Dict[str, Any]:
|
||||
"""
|
||||
获取 Excel 表格元数据(供 LLM 摘要使用)
|
||||
|
||||
Args:
|
||||
filepath: Excel 文件路径
|
||||
sheet_name: Sheet 名称(可选,默认第一个 sheet)
|
||||
|
||||
Returns:
|
||||
{
|
||||
'sheet_name': str,
|
||||
'columns': List[str],
|
||||
'row_count': int,
|
||||
'col_count': int,
|
||||
'sample_rows': List[Dict], # 前 5 行数据
|
||||
}
|
||||
"""
|
||||
filepath = Path(filepath)
|
||||
if not filepath.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {filepath}")
|
||||
|
||||
xls = pd.ExcelFile(filepath)
|
||||
|
||||
if sheet_name is None:
|
||||
sheet_name = xls.sheet_names[0]
|
||||
|
||||
df = pd.read_excel(filepath, sheet_name=sheet_name)
|
||||
df = df.fillna('')
|
||||
|
||||
columns = [str(col) for col in df.columns.tolist()]
|
||||
row_count = len(df)
|
||||
col_count = len(df.columns)
|
||||
|
||||
# 前 5 行样本
|
||||
sample_df = df.head(5)
|
||||
sample_rows = sample_df.to_dict(orient='records')
|
||||
|
||||
return {
|
||||
'sheet_name': sheet_name,
|
||||
'columns': columns,
|
||||
'row_count': row_count,
|
||||
'col_count': col_count,
|
||||
'sample_rows': sample_rows
|
||||
}
|
||||
|
||||
|
||||
def convert_to_rag_format(result: Dict[str, Any]) -> List[Dict]:
|
||||
"""
|
||||
将 Excel 解析结果转换为 RAG 入库格式
|
||||
|
||||
Args:
|
||||
result: parse_excel() 返回结果
|
||||
|
||||
Returns:
|
||||
[{'text': ..., 'page': ..., 'has_table': True, ...}, ...]
|
||||
"""
|
||||
pages_content = []
|
||||
|
||||
for chunk in result['chunks']:
|
||||
# 构建内容文本
|
||||
content = f"【表格】{chunk.title}\n\n{chunk.content}"
|
||||
|
||||
page_info = {
|
||||
'text': content,
|
||||
'page': chunk.page_start,
|
||||
'page_end': chunk.page_end,
|
||||
'has_table': True,
|
||||
'section': chunk.title,
|
||||
'section_path': chunk.section_path,
|
||||
'level': 0,
|
||||
'chunk_type': 'table',
|
||||
'source_file': chunk.source_file,
|
||||
'is_excel_chunk': True, # 标记为 Excel 输出
|
||||
# Excel 专用元数据
|
||||
'sheet_name': chunk.sheet_name,
|
||||
'row_start': chunk.row_start,
|
||||
'row_end': chunk.row_end,
|
||||
'col_count': chunk.col_count,
|
||||
}
|
||||
|
||||
pages_content.append(page_info)
|
||||
|
||||
return pages_content
|
||||
|
||||
|
||||
# ========== 兼容旧接口 ==========
|
||||
|
||||
def parse_xlsx_enhanced(filepath: str) -> Dict[str, Any]:
|
||||
"""
|
||||
兼容旧接口:使用增强解析器处理 Excel 文件
|
||||
|
||||
Args:
|
||||
filepath: Excel 文件路径
|
||||
|
||||
Returns:
|
||||
解析结果(兼容旧格式)
|
||||
"""
|
||||
result = parse_excel(filepath)
|
||||
|
||||
# 转换为旧格式
|
||||
chunks = []
|
||||
for chunk in result['chunks']:
|
||||
chunks.append({
|
||||
'content': chunk.content,
|
||||
'title': chunk.title,
|
||||
'sheet': chunk.sheet_name,
|
||||
'row_range': f"{chunk.row_start+1}-{chunk.row_end}",
|
||||
'col_range': f"A-{chr(64+chunk.col_count)}" if chunk.col_count <= 26 else "A-...",
|
||||
'chunk_type': chunk.chunk_type,
|
||||
'headers': chunk.headers,
|
||||
'source_file': chunk.source_file,
|
||||
'metadata': {
|
||||
'row_count': chunk.row_end - chunk.row_start,
|
||||
'col_count': chunk.col_count
|
||||
}
|
||||
})
|
||||
|
||||
return {
|
||||
'chunks': chunks,
|
||||
'sheets': [{'name': s, 'rows': 0, 'cols': 0} for s in result['sheets']],
|
||||
'metadata': {
|
||||
'source_file': result['source_file'],
|
||||
'total_chunks': len(chunks)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def get_excel_chunks_for_rag(
|
||||
filepath: str,
|
||||
min_chunk_size: int = 50
|
||||
) -> tuple:
|
||||
"""
|
||||
兼容旧接口:获取适合 RAG 系统的 Excel 分块
|
||||
|
||||
Args:
|
||||
filepath: 文件路径
|
||||
min_chunk_size: 最小分块大小
|
||||
|
||||
Returns:
|
||||
(documents, metadatas) - 文档列表和元数据列表
|
||||
"""
|
||||
result = parse_excel(filepath)
|
||||
|
||||
documents = []
|
||||
metadatas = []
|
||||
|
||||
for chunk in result['chunks']:
|
||||
if len(chunk.content.strip()) >= min_chunk_size:
|
||||
documents.append(chunk.content)
|
||||
metadatas.append({
|
||||
'title': chunk.title,
|
||||
'sheet': chunk.sheet_name,
|
||||
'row_range': f"{chunk.row_start+1}-{chunk.row_end}",
|
||||
'col_count': chunk.col_count,
|
||||
'chunk_type': chunk.chunk_type,
|
||||
'source_file': chunk.source_file
|
||||
})
|
||||
|
||||
return documents, metadatas
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
print("用法: python excel_parser.py <Excel文件路径>")
|
||||
sys.exit(1)
|
||||
|
||||
file_path = sys.argv[1]
|
||||
|
||||
print(f"正在解析: {file_path}")
|
||||
result = parse_excel(file_path)
|
||||
|
||||
print(f"\n解析完成:")
|
||||
print(f"- Sheets: {result['sheets']}")
|
||||
print(f"- 总行数: {result['total_rows']}")
|
||||
print(f"- 表格块数: {len(result['chunks'])}")
|
||||
|
||||
# 显示每个块的信息
|
||||
print("\n表格块详情:")
|
||||
for i, chunk in enumerate(result['chunks']):
|
||||
print(f"\n--- Chunk {i+1} ---")
|
||||
print(f"Sheet: {chunk.sheet_name}")
|
||||
print(f"行范围: {chunk.row_start+1} - {chunk.row_end}")
|
||||
print(f"列数: {chunk.col_count}")
|
||||
preview = chunk.content[:200] + "..." if len(chunk.content) > 200 else chunk.content
|
||||
print(f"内容预览: {preview}")
|
||||
495
parsers/image_extractor.py
Normal file
495
parsers/image_extractor.py
Normal file
@@ -0,0 +1,495 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
PDF 图片提取模块
|
||||
|
||||
从 PDF 文档中提取图片,支持:
|
||||
- 使用 PyMuPDF (fitz) 提取嵌入式图片
|
||||
- 保存到指定目录
|
||||
- 生成图片元数据供 RAG 系统使用
|
||||
"""
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional, Tuple, Any
|
||||
from dataclasses import dataclass, asdict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageInfo:
|
||||
"""图片信息"""
|
||||
image_id: str # 图片唯一 ID
|
||||
original_name: str # 原始文件名
|
||||
storage_path: str # 存储路径(相对路径)
|
||||
page: int # 所在页码
|
||||
width: int # 宽度
|
||||
height: int # 高度
|
||||
format: str # 格式 (png, jpg, etc.)
|
||||
size_bytes: int # 文件大小
|
||||
caption: str = "" # 图片说明(可选)
|
||||
bbox: Optional[List[float]] = None # 边界框坐标
|
||||
|
||||
|
||||
def extract_images_from_pdf(
|
||||
pdf_path: str,
|
||||
output_dir: str,
|
||||
min_width: int = 100,
|
||||
min_height: int = 100,
|
||||
max_width: int = 2000, # 新增:最大宽度阈值(过滤跨页底纹)
|
||||
max_height: int = 2000, # 新增:最大高度阈值
|
||||
max_images: int = 50
|
||||
) -> List[ImageInfo]:
|
||||
"""
|
||||
从 PDF 中提取图片
|
||||
|
||||
Args:
|
||||
pdf_path: PDF 文件路径
|
||||
output_dir: 图片输出目录
|
||||
min_width: 最小宽度阈值(过滤小图标)
|
||||
min_height: 最小高度阈值
|
||||
max_width: 最大宽度阈值(过滤跨页底纹、背景横幅)
|
||||
max_height: 最大高度阈值
|
||||
max_images: 最大提取图片数量
|
||||
|
||||
Returns:
|
||||
图片信息列表
|
||||
"""
|
||||
try:
|
||||
import fitz # PyMuPDF
|
||||
except ImportError:
|
||||
logger.warning("PyMuPDF 未安装,无法提取图片。请运行: pip install PyMuPDF")
|
||||
return []
|
||||
|
||||
pdf_path = Path(pdf_path)
|
||||
if not pdf_path.exists():
|
||||
raise FileNotFoundError(f"PDF 文件不存在: {pdf_path}")
|
||||
|
||||
# 创建输出目录
|
||||
output_dir = Path(output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
images = []
|
||||
image_count = 0
|
||||
|
||||
try:
|
||||
doc = fitz.open(str(pdf_path))
|
||||
filename = pdf_path.stem
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
image_list = page.get_images(full=True)
|
||||
|
||||
for img_index, img_info in enumerate(image_list):
|
||||
if image_count >= max_images:
|
||||
break
|
||||
|
||||
try:
|
||||
# 获取图片引用
|
||||
xref = img_info[0]
|
||||
|
||||
# 提取图片
|
||||
base_image = doc.extract_image(xref)
|
||||
if not base_image:
|
||||
continue
|
||||
|
||||
image_bytes = base_image.get("image")
|
||||
if not image_bytes:
|
||||
continue
|
||||
|
||||
# 获取图片属性
|
||||
width = base_image.get("width", 0)
|
||||
height = base_image.get("height", 0)
|
||||
image_ext = base_image.get("ext", "png")
|
||||
|
||||
# 过滤太小的图片(通常是图标)
|
||||
if width < min_width or height < min_height:
|
||||
continue
|
||||
|
||||
# 过滤太大的图片(跨页底纹、背景横幅)
|
||||
if width > max_width or height > max_height:
|
||||
continue
|
||||
|
||||
# 生成唯一 ID
|
||||
image_id = f"{filename}_p{page_num + 1}_img{img_index + 1}"
|
||||
|
||||
# 确定文件扩展名
|
||||
if image_ext in ["jpeg", "jpg"]:
|
||||
ext = ".jpg"
|
||||
elif image_ext == "png":
|
||||
ext = ".png"
|
||||
else:
|
||||
ext = f".{image_ext}"
|
||||
|
||||
# 保存图片
|
||||
image_filename = f"{image_id}{ext}"
|
||||
image_path = output_dir / image_filename
|
||||
|
||||
with open(image_path, "wb") as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
# 记录图片信息
|
||||
images.append(ImageInfo(
|
||||
image_id=image_id,
|
||||
original_name=filename,
|
||||
storage_path=f"images/{image_filename}",
|
||||
page=page_num + 1,
|
||||
width=width,
|
||||
height=height,
|
||||
format=image_ext,
|
||||
size_bytes=len(image_bytes),
|
||||
caption=f"图片 {img_index + 1}"
|
||||
))
|
||||
|
||||
image_count += 1
|
||||
|
||||
except Exception as e:
|
||||
# 单个图片提取失败不影响其他图片
|
||||
logger.warning(f"提取图片失败 (页 {page_num + 1}, 图片 {img_index + 1}): {e}")
|
||||
continue
|
||||
|
||||
if image_count >= max_images:
|
||||
break
|
||||
|
||||
doc.close()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"PDF 图片提取失败: {e}")
|
||||
return []
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def extract_images_batch(
|
||||
pdf_dir: str,
|
||||
output_dir: str,
|
||||
**kwargs
|
||||
) -> Dict[str, List[ImageInfo]]:
|
||||
"""
|
||||
批量提取 PDF 目录下所有文件的图片
|
||||
|
||||
Args:
|
||||
pdf_dir: PDF 文件目录
|
||||
output_dir: 图片输出目录
|
||||
**kwargs: 传递给 extract_images_from_pdf 的参数
|
||||
|
||||
Returns:
|
||||
{文件名: [ImageInfo, ...], ...}
|
||||
"""
|
||||
pdf_dir = Path(pdf_dir)
|
||||
output_dir = Path(output_dir)
|
||||
|
||||
results = {}
|
||||
|
||||
for pdf_file in pdf_dir.glob("**/*.pdf"):
|
||||
try:
|
||||
# 为每个 PDF 创建子目录
|
||||
pdf_output_dir = output_dir / pdf_file.stem
|
||||
|
||||
images = extract_images_from_pdf(
|
||||
str(pdf_file),
|
||||
str(pdf_output_dir),
|
||||
**kwargs
|
||||
)
|
||||
|
||||
if images:
|
||||
results[pdf_file.name] = images
|
||||
logger.info(f"{pdf_file.name}: 提取 {len(images)} 张图片")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"{pdf_file.name}: {e}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def get_images_base_path() -> str:
|
||||
"""获取图片存储的基础路径"""
|
||||
try:
|
||||
from config import DOCUMENTS_PATH
|
||||
return os.path.join(DOCUMENTS_PATH, "images")
|
||||
except ImportError:
|
||||
return "documents/images"
|
||||
|
||||
|
||||
# ==================== 集成到现有解析器 ====================
|
||||
|
||||
def filter_noise_images(
|
||||
images: List[Any],
|
||||
min_size: int = 100,
|
||||
max_size: int = 2000,
|
||||
enable_hash_dedup: bool = True,
|
||||
enable_content_check: bool = True,
|
||||
max_aspect_ratio: float = 10.0
|
||||
) -> List[Any]:
|
||||
"""
|
||||
三级噪音图片过滤管道
|
||||
|
||||
Level 1: 尺寸过滤(过滤图标、跨页底纹)
|
||||
Level 2: Hash 去重(过滤重复图片)
|
||||
Level 3: 内容检测(过滤纯色背景、装饰横幅)
|
||||
|
||||
Args:
|
||||
images: 图片列表(ImageInfo 或 dict)
|
||||
min_size: 最小尺寸阈值(像素),过滤图标
|
||||
max_size: 最大尺寸阈值(像素),过滤跨页底纹
|
||||
enable_hash_dedup: 启用 Hash 去重
|
||||
enable_content_check: 启用内容相关性检测
|
||||
max_aspect_ratio: 最大宽高比阈值,过滤装饰横幅
|
||||
|
||||
Returns:
|
||||
过滤后的图片列表
|
||||
"""
|
||||
if not images:
|
||||
return images
|
||||
|
||||
filtered = []
|
||||
seen_hashes = set()
|
||||
|
||||
for img in images:
|
||||
# 支持 dataclass 和 dict 两种格式
|
||||
if hasattr(img, 'width'):
|
||||
width, height = img.width, img.height
|
||||
storage_path = getattr(img, 'storage_path', '')
|
||||
else:
|
||||
width = img.get('width', 0)
|
||||
height = img.get('height', 0)
|
||||
storage_path = img.get('storage_path', '')
|
||||
|
||||
# Level 1: 尺寸过滤
|
||||
if width < min_size or height < min_size:
|
||||
continue # 图标、装饰线条
|
||||
if width > max_size or height > max_size:
|
||||
continue # 跨页底纹、背景横幅
|
||||
|
||||
# Level 2: Hash 去重
|
||||
if enable_hash_dedup and storage_path:
|
||||
try:
|
||||
img_hash = _compute_image_hash(storage_path)
|
||||
if img_hash in seen_hashes:
|
||||
continue # 重复图片
|
||||
seen_hashes.add(img_hash)
|
||||
except Exception:
|
||||
pass # Hash 计算失败时跳过去重
|
||||
|
||||
# Level 3: 内容相关性检测
|
||||
if enable_content_check:
|
||||
aspect_ratio = max(width, height) / max(min(width, height), 1)
|
||||
|
||||
# 过滤极端宽高比(装饰横幅)
|
||||
if aspect_ratio > max_aspect_ratio:
|
||||
continue
|
||||
|
||||
# 过滤纯色/渐变背景
|
||||
if storage_path and _is_solid_color_image(storage_path):
|
||||
continue
|
||||
|
||||
filtered.append(img)
|
||||
|
||||
return filtered
|
||||
|
||||
|
||||
def _compute_image_hash(image_path: str) -> str:
|
||||
"""
|
||||
计算图片文件的 Hash 值
|
||||
|
||||
Args:
|
||||
image_path: 图片文件路径
|
||||
|
||||
Returns:
|
||||
MD5 Hash 字符串
|
||||
"""
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(image_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b''):
|
||||
hash_md5.update(chunk)
|
||||
return hash_md5.hexdigest()
|
||||
|
||||
|
||||
def _is_solid_color_image(image_path: str, threshold: float = 0.95) -> bool:
|
||||
"""
|
||||
检测图片是否为纯色/渐变背景(装饰性横幅)
|
||||
|
||||
使用简单的颜色分布检测:
|
||||
- 如果图片 95% 以上像素属于同一颜色范围,判定为纯色背景
|
||||
|
||||
Args:
|
||||
image_path: 图片文件路径
|
||||
threshold: 纯色判定阈值
|
||||
|
||||
Returns:
|
||||
True 表示是纯色背景图片
|
||||
"""
|
||||
try:
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
except ImportError:
|
||||
# PIL/numpy 未安装,跳过检测
|
||||
return False
|
||||
|
||||
try:
|
||||
img = Image.open(image_path)
|
||||
# 缩小图片加速处理
|
||||
img.thumbnail((100, 100))
|
||||
img_array = np.array(img)
|
||||
|
||||
if img_array.ndim == 2:
|
||||
# 灰度图
|
||||
unique, counts = np.unique(img_array, return_counts=True)
|
||||
elif img_array.ndim == 3:
|
||||
# 彩色图,计算颜色直方图
|
||||
pixels = img_array.reshape(-1, img_array.shape[-1])
|
||||
# 量化颜色(减少颜色数量)
|
||||
quantized = (pixels // 32) * 32
|
||||
unique, counts = np.unique(quantized, axis=0, return_counts=True)
|
||||
else:
|
||||
return False
|
||||
|
||||
# 如果主颜色占比超过阈值,判定为纯色背景
|
||||
max_color_ratio = max(counts) / sum(counts)
|
||||
return max_color_ratio > threshold
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def enrich_chunks_with_images(
|
||||
chunks: List[Any],
|
||||
images: List[ImageInfo],
|
||||
source_file: str
|
||||
) -> List[Any]:
|
||||
"""
|
||||
为分块添加图片信息
|
||||
|
||||
根据页码将图片关联到对应的分块
|
||||
|
||||
Args:
|
||||
chunks: 分块列表(ChunkMetadata 或 dict)
|
||||
images: 图片信息列表
|
||||
source_file: 源文件名
|
||||
|
||||
Returns:
|
||||
添加了图片信息的分块列表
|
||||
"""
|
||||
if not images:
|
||||
return chunks
|
||||
|
||||
# 按页码分组图片
|
||||
page_to_images = {}
|
||||
for img in images:
|
||||
page = img.page
|
||||
if page not in page_to_images:
|
||||
page_to_images[page] = []
|
||||
page_to_images[page].append({
|
||||
"id": img.image_id,
|
||||
"caption": img.caption,
|
||||
"page": img.page,
|
||||
"width": img.width,
|
||||
"height": img.height
|
||||
})
|
||||
|
||||
# 为每个分块添加图片信息
|
||||
for chunk in chunks:
|
||||
# 支持 dataclass 和 dict 两种格式
|
||||
if hasattr(chunk, 'page_start'):
|
||||
page_start = chunk.page_start
|
||||
page_end = getattr(chunk, 'page_end', page_start)
|
||||
else:
|
||||
page_start = chunk.get('page_start', 1)
|
||||
page_end = chunk.get('page_end', page_start)
|
||||
|
||||
# 单页绑定:仅在切片不跨页时绑定图片
|
||||
chunk_images = []
|
||||
if page_start == page_end and page_start in page_to_images:
|
||||
chunk_images = page_to_images[page_start]
|
||||
|
||||
# 应用噪音过滤
|
||||
chunk_images = filter_noise_images(chunk_images)
|
||||
|
||||
# 限制每切片最多 3 张图片
|
||||
chunk_images = chunk_images[:3]
|
||||
|
||||
# 添加到分块
|
||||
if chunk_images:
|
||||
if hasattr(chunk, '__dict__'):
|
||||
# dataclass
|
||||
chunk.images = chunk_images
|
||||
else:
|
||||
# dict
|
||||
chunk['images'] = chunk_images
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
# ==================== 测试 ====================
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
print("=" * 60)
|
||||
print("PDF 图片提取模块测试")
|
||||
print("=" * 60)
|
||||
|
||||
# 检查依赖
|
||||
try:
|
||||
import fitz
|
||||
print("[OK] PyMuPDF 已安装")
|
||||
except ImportError:
|
||||
print("[错误] PyMuPDF 未安装,请运行: pip install PyMuPDF")
|
||||
sys.exit(1)
|
||||
|
||||
# 测试提取
|
||||
if len(sys.argv) >= 2:
|
||||
pdf_path = sys.argv[1]
|
||||
output_dir = sys.argv[2] if len(sys.argv) >= 3 else "documents/images"
|
||||
|
||||
print(f"\n提取图片: {pdf_path}")
|
||||
print(f"输出目录: {output_dir}")
|
||||
|
||||
images = extract_images_from_pdf(pdf_path, output_dir)
|
||||
|
||||
print(f"\n提取结果: {len(images)} 张图片")
|
||||
for img in images[:10]:
|
||||
print(f" - {img.image_id}: {img.width}x{img.height}, {img.size_bytes} bytes, 页码 {img.page}")
|
||||
else:
|
||||
print("\n用法: python image_extractor.py <pdf_path> [output_dir]")
|
||||
print("\n功能演示: 创建模拟图片信息")
|
||||
|
||||
# 创建模拟数据演示功能
|
||||
mock_images = [
|
||||
ImageInfo(
|
||||
image_id="test_p1_img1",
|
||||
original_name="test.pdf",
|
||||
storage_path="images/test_p1_img1.png",
|
||||
page=1,
|
||||
width=800,
|
||||
height=600,
|
||||
format="png",
|
||||
size_bytes=45000,
|
||||
caption="流程图"
|
||||
),
|
||||
ImageInfo(
|
||||
image_id="test_p3_img1",
|
||||
original_name="test.pdf",
|
||||
storage_path="images/test_p3_img1.jpg",
|
||||
page=3,
|
||||
width=1200,
|
||||
height=900,
|
||||
format="jpg",
|
||||
size_bytes=120000,
|
||||
caption="组织架构图"
|
||||
)
|
||||
]
|
||||
|
||||
print("\n模拟图片信息:")
|
||||
for img in mock_images:
|
||||
print(f" ID: {img.image_id}")
|
||||
print(f" 页码: {img.page}")
|
||||
print(f" 尺寸: {img.width}x{img.height}")
|
||||
print(f" 格式: {img.format}")
|
||||
print(f" 大小: {img.size_bytes} bytes")
|
||||
print()
|
||||
1678
parsers/mineru_parser.py
Normal file
1678
parsers/mineru_parser.py
Normal file
File diff suppressed because it is too large
Load Diff
31
parsers/pdf_mineru.py
Normal file
31
parsers/pdf_mineru.py
Normal file
@@ -0,0 +1,31 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
MinerU PDF 解析模块(兼容性别名)
|
||||
|
||||
此文件保留用于向后兼容,实际实现已迁移到 mineru_parser.py
|
||||
新代码请直接使用: from parsers.mineru_parser import parse_with_mineru
|
||||
"""
|
||||
|
||||
# 从新模块导入所有内容
|
||||
from parsers.mineru_parser import (
|
||||
parse_with_mineru,
|
||||
parse_pdf_with_mineru,
|
||||
parse_docx_with_mineru,
|
||||
parse_xlsx_with_mineru,
|
||||
parse_pptx_with_mineru,
|
||||
convert_to_rag_format,
|
||||
html_table_to_markdown,
|
||||
MinerUChunk,
|
||||
ChunkMetadata,
|
||||
SUPPORTED_FORMATS,
|
||||
)
|
||||
|
||||
# 保持向后兼容
|
||||
__all__ = [
|
||||
'parse_pdf_with_mineru',
|
||||
'parse_with_mineru',
|
||||
'convert_to_rag_format',
|
||||
'html_table_to_markdown',
|
||||
'MinerUChunk',
|
||||
'ChunkMetadata',
|
||||
]
|
||||
27
parsers/txt_parser.py
Normal file
27
parsers/txt_parser.py
Normal file
@@ -0,0 +1,27 @@
|
||||
"""
|
||||
TXT 文本解析器
|
||||
|
||||
简单的文本文件读取,支持 UTF-8 和 GBK 编码自动检测。
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def extract_text_from_txt(filepath):
|
||||
"""从TXT提取文本"""
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
return f.read()
|
||||
except UnicodeDecodeError:
|
||||
# 尝试其他编码
|
||||
try:
|
||||
with open(filepath, 'r', encoding='gbk') as f:
|
||||
return f.read()
|
||||
except Exception as e:
|
||||
logger.error(f"TXT解析错误 {filepath}: {e}")
|
||||
return ""
|
||||
except Exception as e:
|
||||
print(f" TXT解析错误 {filepath}: {e}")
|
||||
return ""
|
||||
Reference in New Issue
Block a user