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