pengxiao
37c0364e7e
Replace all print()/stderr logging with Python logging module using logger = logging.getLogger(__name__) pattern for consistent log levels and formatting. Extract score conversion to shared score_utils module. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
"""
|
|
LangChain 1.0 文档加载器封装 + PDF图片提取
|
|
"""
|
|
import logging
|
|
from typing import List, Optional, Dict
|
|
from pathlib import Path
|
|
import fitz # pymupdf
|
|
from langchain_community.document_loaders import (
|
|
PyPDFLoader,
|
|
Docx2txtLoader,
|
|
TextLoader,
|
|
UnstructuredMarkdownLoader
|
|
)
|
|
from langchain_core.documents import Document
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PDFImageExtractor:
|
|
"""使用pymupdf从PDF中提取内嵌图片"""
|
|
|
|
@staticmethod
|
|
def extract_images(file_path: str, output_dir: str) -> List[dict]:
|
|
"""提取PDF中所有图片，返回图片元数据列表"""
|
|
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
|
images = []
|
|
doc = fitz.open(file_path)
|
|
|
|
for page_num in range(len(doc)):
|
|
page = doc[page_num]
|
|
# 获取页面文本作为图片上下文
|
|
page_text = page.get_text("text")
|
|
# 提取页面内嵌图片
|
|
image_list = page.get_images(full=True)
|
|
|
|
for img_idx, img_info in enumerate(image_list):
|
|
xref = img_info[0]
|
|
try:
|
|
base_image = doc.extract_image(xref)
|
|
image_bytes = base_image["image"]
|
|
ext = base_image["ext"]
|
|
|
|
# 过滤太小的图片（图标、装饰元素等）
|
|
if len(image_bytes) < 2048:
|
|
continue
|
|
|
|
# 图片周围文本（取该页文字前后各300字作为上下文）
|
|
context_start = max(0, page_text.find(
|
|
page_text[:len(page_text)//2]) if len(page_text) > 600
|
|
else 0)
|
|
context_text = page_text[context_start:context_start+600].strip()
|
|
|
|
filename = f"page{page_num+1}_img{img_idx+1}.{ext}"
|
|
output_path = Path(output_dir) / filename
|
|
output_path.write_bytes(image_bytes)
|
|
|
|
images.append({
|
|
"path": str(output_path),
|
|
"filename": filename,
|
|
"page": page_num + 1,
|
|
"context_text": context_text,
|
|
"size": len(image_bytes),
|
|
})
|
|
except Exception as e:
|
|
logger.warning(f"提取图片失败 page={page_num+1} img={img_idx}: {e}")
|
|
continue
|
|
|
|
doc.close()
|
|
return images
|
|
|
|
|
|
class DocumentLoaderFactory:
|
|
"""文档加载器工厂"""
|
|
|
|
@staticmethod
|
|
def get_loader(file_path: str, file_type: str):
|
|
"""根据文件类型获取对应的加载器"""
|
|
loaders = {
|
|
".pdf": PyPDFLoader,
|
|
".docx": Docx2txtLoader,
|
|
".txt": TextLoader,
|
|
".md": UnstructuredMarkdownLoader,
|
|
}
|
|
|
|
loader_class = loaders.get(file_type)
|
|
if not loader_class:
|
|
raise ValueError(f"Unsupported file type: {file_type}")
|
|
|
|
return loader_class(file_path)
|
|
|
|
@staticmethod
|
|
def load_document(file_path: str, file_type: str, metadata: Optional[dict] = None) -> List[Document]:
|
|
"""加载文档并添加元数据"""
|
|
loader = DocumentLoaderFactory.get_loader(file_path, file_type)
|
|
documents = loader.load()
|
|
|
|
if metadata:
|
|
for doc in documents:
|
|
doc.metadata.update(metadata)
|
|
|
|
return documents
|