101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
演示优化后的图片OCR与RAG系统集成
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import asyncio
|
||
from pathlib import Path
|
||
|
||
# 添加源码路径
|
||
sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
|
||
|
||
from base_rag.core import BaseRAG
|
||
from base_rag.image_processor import ImageProcessor
|
||
|
||
|
||
class DemoRAG(BaseRAG):
|
||
"""演示RAG实现"""
|
||
|
||
async def ingest(self, file_path: str, **kwargs):
|
||
"""文档导入"""
|
||
return await self.process_file_to_vector_store(file_path, **kwargs)
|
||
|
||
async def query(self, question: str) -> str:
|
||
"""查询实现"""
|
||
docs = await self.similarity_search_with_rerank(question, k=3)
|
||
|
||
if not docs:
|
||
return "抱歉,没有找到相关信息。"
|
||
|
||
# 简单的结果组织
|
||
context = "\n".join([doc.page_content for doc in docs])
|
||
return f"基于以下内容回答:\n{context}"
|
||
|
||
|
||
async def demo_image_ocr_integration():
|
||
"""演示图片OCR与RAG系统集成"""
|
||
print("🎯 演示:优化后的图片OCR与RAG系统集成")
|
||
print("=" * 60)
|
||
|
||
# 配置RAG系统,启用图片处理
|
||
image_config = {
|
||
"enabled": True,
|
||
"type": "local", # 使用本地模式(BLIP + EasyOCR)
|
||
"engine": "easyocr"
|
||
}
|
||
|
||
try:
|
||
# 初始化RAG系统
|
||
print("🚀 初始化RAG系统...")
|
||
rag = DemoRAG(
|
||
persist_directory="./demo_chroma_ocr",
|
||
image_config=image_config
|
||
)
|
||
|
||
# 检查是否有图片文件需要处理
|
||
image_files = []
|
||
test_dirs = ["./examples/", "./demo_documents/", "./"]
|
||
|
||
for test_dir in test_dirs:
|
||
if os.path.exists(test_dir):
|
||
for file in os.listdir(test_dir):
|
||
if file.lower().endswith(('.png', '.jpg', '.jpeg')):
|
||
image_files.append(os.path.join(test_dir, file))
|
||
|
||
if image_files:
|
||
print(f"\n📷 发现 {len(image_files)} 个图片文件")
|
||
|
||
# 处理图片文件
|
||
for img_file in image_files[:2]: # 限制处理数量
|
||
print(f"\n🔍 处理图片: {os.path.basename(img_file)}")
|
||
try:
|
||
# 直接测试图片处理器
|
||
processor = ImageProcessor(image_config)
|
||
result = processor.extract_image_description(img_file)
|
||
print(f"📝 OCR结果:\n{result[:200]}...")
|
||
|
||
# 这里可以将图片内容添加到向量库
|
||
# await rag.ingest(img_file)
|
||
|
||
except Exception as e:
|
||
print(f"❌ 处理失败: {e}")
|
||
else:
|
||
print("⚠️ 未找到测试图片")
|
||
|
||
print(f"\n✅ 演示完成!")
|
||
print("\n🌟 优化亮点:")
|
||
print(" • 使用EasyOCR进行高质量文字识别")
|
||
print(" • local模式结合图片描述和OCR文本")
|
||
print(" • api模式也会自动加入OCR文本内容")
|
||
print(" • basic模式专注于OCR文字提取")
|
||
print(" • 所有模式都支持中英文混合识别")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 演示失败: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(demo_image_ocr_integration())
|