base_rag/demo_image_integration.py

101 lines
3.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/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())