82 lines
2.0 KiB
Python
82 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""测试 PDF 多模态提取完整链路
|
|
|
|
用法:
|
|
python scripts/test_multimodal.py
|
|
"""
|
|
|
|
import io
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Windows 终端强制 UTF-8
|
|
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT)) # noqa: E402
|
|
|
|
from src.doc.llm_extractor import extract_document # noqa: E402
|
|
from src.doc.pdf import render_pdf_to_images # noqa: E402
|
|
|
|
|
|
def test_render() -> None:
|
|
"""测试 PDF 渲染"""
|
|
pdf_path = ROOT / "事前申请单.pdf"
|
|
if not pdf_path.exists():
|
|
print(f"跳过: {pdf_path.name} 不存在")
|
|
return
|
|
|
|
print("=" * 60)
|
|
print("测试 PDF 渲染")
|
|
print("=" * 60)
|
|
try:
|
|
images = render_pdf_to_images(pdf_path, dpi=150)
|
|
if images:
|
|
print(f"成功渲染 {len(images)} 页")
|
|
for i, img_b64 in enumerate(images):
|
|
print(f" 第 {i + 1} 页: base64 长度 {len(img_b64)} 字符")
|
|
else:
|
|
print("渲染返回空列表！")
|
|
except ImportError as e:
|
|
print(f"导入失败: {e}")
|
|
except Exception as e:
|
|
print(f"渲染异常: {e}")
|
|
|
|
|
|
def test_multimodal_extract() -> None:
|
|
"""测试完整的多模态提取链路"""
|
|
pdf_path = ROOT / "事前申请单.pdf"
|
|
if not pdf_path.exists():
|
|
print(f"跳过: {pdf_path.name} 不存在")
|
|
return
|
|
|
|
print()
|
|
print("=" * 60)
|
|
print("测试多模态提取")
|
|
print("=" * 60)
|
|
try:
|
|
result = extract_document(pdf_path)
|
|
if result:
|
|
print("提取成功:")
|
|
import json
|
|
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
else:
|
|
print("提取返回空字典！")
|
|
except Exception as e:
|
|
print(f"提取异常: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
|
|
|
|
def main() -> None:
|
|
test_render()
|
|
test_multimodal_extract()
|
|
print()
|
|
print("测试完成!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|