feat: 删除多余文件

This commit is contained in:
李如威 2025-08-09 10:33:44 +08:00
parent 3e4df252da
commit 51c83ff4cc
15 changed files with 54 additions and 831 deletions

View File

@ -1,117 +0,0 @@
# 项目目录整理总结
## 🗂️ 整理前的问题
- 存在多个重复的测试目录:`test_files/`, `test_docs/`, `documents/`
- 测试文件分散在不同位置,管理混乱
- 大量临时测试脚本和数据库文件
- 重复的测试文件和配置
## 🧹 清理操作
### 删除的目录和文件
1. **重复目录**:
- `test_docs/` - 删除
- `documents/` - 删除
2. **临时测试脚本**:
- `create_docx.py`
- `test_docx.py`
- `test_file_formats.py`
- `comprehensive_test.py`
- `final_test.py`
3. **多余数据库文件**:
- `demo_file_status.db`
- `file_status.db`
- `test_status.db`
- `final_test.db`
4. **临时文档**:
- `CLEANUP_SUMMARY.md`
- `TEST_REPORT.md`
5. **重复测试文件**:
- `test_files/test_document.txt`
- `test_files/test_markdown.md`
## 📁 整理后的目录结构
```
base_rag/
├── src/ # 源代码
│ └── base_rag/
│ ├── __init__.py
│ └── core.py
├── examples/ # 示例代码
│ ├── simple_test.py # 基础功能测试
│ └── multi_format_test.py # 多格式文件测试
├── test_files/ # 统一的测试文件目录
│ ├── knowledge.txt # TXT格式测试文件
│ ├── python_basics.txt # Python基础知识
│ ├── web_frameworks.txt # Web框架知识
│ ├── data_science.txt # 数据科学知识
│ ├── python_guide.md # Python指南(MD格式)
│ ├── machine_learning.md # 机器学习(MD格式)
│ └── deep_learning_guide.docx # 深度学习(DOCX格式)
├── chroma_db/ # 向量数据库
├── scripts/ # 构建脚本
├── venv/ # Python虚拟环境
├── README.md # 项目说明
├── requirements.txt # 依赖包
└── pyproject.toml # 项目配置
```
## 🎯 统一的测试配置
### 测试目录统一为 `test_files/`
- 所有测试文件集中在 `test_files/` 目录
- 包含 TXT、MD、DOCX 三种格式的测试文件
- 涵盖不同主题Python、机器学习、深度学习、数据科学
### 数据库统一为 `status.db`
- 使用单一的状态数据库文件
- 避免多个测试脚本创建重复的数据库
### 示例脚本优化
1. **`simple_test.py`** - 基础功能测试
- 测试基本的文件处理和查询功能
- 使用 TXT 格式文件进行测试
2. **`multi_format_test.py`** - 多格式文件测试
- 测试 TXT、MD、DOCX 三种格式
- 验证跨格式查询功能
- 展示异步处理能力
## ✅ 整理效果
1. **目录结构清晰**: 消除了重复目录,统一了测试文件位置
2. **文件管理规范**: 删除了临时和重复文件
3. **测试配置统一**: 所有测试使用相同的目录和数据库配置
4. **代码组织优化**: 保留了两个核心的测试示例
## 🚀 后续测试指南
### 运行基础测试
```bash
cd /Users/liruwei/Documents/code/project/demo/base_rag
python examples/simple_test.py
```
### 运行多格式测试
```bash
cd /Users/liruwei/Documents/code/project/demo/base_rag
python examples/multi_format_test.py
```
### 添加新测试文件
将新的测试文件直接放入 `test_files/` 目录即可,支持的格式:
- `.txt` - 纯文本文件
- `.md` - Markdown文件
- `.docx` - Word文档
现在项目结构更加清晰,便于后续的开发和测试!

View File

@ -1,170 +0,0 @@
# 文件处理功能说明
## 概述
BaseRAG 类现在支持自动处理多种格式的文件,并将它们转换为向量嵌入存储到知识库中。该功能包括:
1. **支持的文件格式**txt、md、doc/docx
2. **文件状态管理**:使用 SQLite 数据库记录处理状态
3. **文件存储管理**:自动处理文件冲突和覆盖
4. **简洁的API设计**:易于使用和扩展
## 新增的类和功能
### FileStatus 枚举
```python
class FileStatus(Enum):
WAITING = "等待中"
PROCESSING = "处理中"
COMPLETED = "处理完毕"
ERROR = "处理失败"
```
### FileManager 类
负责文件存储和状态管理:
- 文件哈希计算,避免重复处理
- SQLite 数据库记录文件状态
- 自动处理文件名冲突
### BaseRAG 新增方法
#### `await process_file_to_vector_store(file_path, chunk_size=500, chunk_overlap=50)`
主要的文件处理方法:
- 自动检测文件类型
- 保存文件到存储目录
- 切分文档并添加到向量库
- 记录处理状态
#### `await get_file_processing_status(file_hash=None)`
获取文件处理状态:
- 传入 file_hash 获取特定文件状态
- 不传参数获取所有文件状态
#### `await list_files_by_status(status=None)`
按状态筛选文件:
- 传入 FileStatus 枚举获取特定状态的文件
- 不传参数获取所有文件
## 使用示例
### 基本用法
```python
from base_rag.core import BaseRAG, FileStatus
import asyncio
async def main():
# 创建 RAG 实例
rag = SimpleRAG(
vector_store_name="my_knowledge_base",
storage_directory="./documents", # 文件存储目录
status_db_path="./file_status.db" # 状态数据库路径
)
# 处理文件
result = await rag.process_file_to_vector_store("path/to/your/document.txt")
print(result)
# 查看处理状态
status = await rag.get_file_processing_status()
print(status)
asyncio.run(main())
```
### 批量处理文件
```python
import os
from pathlib import Path
import asyncio
async def batch_process():
rag = SimpleRAG()
# 处理目录中的所有文件
docs_dir = Path("./my_documents")
for file_path in docs_dir.glob("*"):
if file_path.suffix.lower() in ['.txt', '.md', '.doc', '.docx']:
print(f"处理文件: {file_path.name}")
result = await rag.process_file_to_vector_store(str(file_path))
print(f"结果: {result['message']}")
# 查看处理结果统计
completed = await rag.list_files_by_status(FileStatus.COMPLETED)
failed = await rag.list_files_by_status(FileStatus.ERROR)
print(f"成功处理: {len(completed)} 个文件")
print(f"处理失败: {len(failed)} 个文件")
asyncio.run(batch_process())
```
## 文件处理流程
1. **文件保存**:计算文件哈希,复制到存储目录,避免重复存储
2. **状态检查**:检查文件是否已经处理过,避免重复处理
3. **状态更新**:将状态更新为"等待中",然后"处理中"
4. **文档加载**:根据文件类型选择适当的加载器
5. **文档切分**:使用 RecursiveCharacterTextSplitter 切分文档
6. **向量化**:将文档片段转换为向量并存储
7. **状态完成**:更新状态为"处理完毕"或"处理失败"
## 数据库结构
SQLite 数据库表结构:
```sql
CREATE TABLE file_status (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT NOT NULL, -- 文件名
file_type TEXT NOT NULL, -- 文件类型 (.txt, .md, .doc 等)
file_hash TEXT UNIQUE NOT NULL, -- 文件哈希值
status TEXT NOT NULL, -- 处理状态
created_at TEXT NOT NULL, -- 创建时间
updated_at TEXT NOT NULL, -- 更新时间
error_message TEXT -- 错误信息(如果有)
);
```
## 配置选项
### BaseRAG 初始化参数
```python
BaseRAG(
vector_store_name="default", # 向量库名称
retriever_top_k=3, # 检索返回文档数
persist_directory="./chroma_db", # 向量库持久化目录
storage_directory="./documents", # 文件存储目录
status_db_path="./file_status.db", # 状态数据库路径
embedding_config=None, # 嵌入模型配置
rerank_config=None, # 重排模型配置
llm=None # 语言模型
)
```
### 文档切分参数
```python
await rag.process_file_to_vector_store(
file_path="document.txt",
chunk_size=500, # 切分块大小
chunk_overlap=50 # 切分重叠大小
)
```
## 安装依赖
确保安装了必要的依赖:
```bash
pip install unstructured python-docx
```
## 注意事项
1. **文件冲突处理**:相同内容的文件会被自动跳过,不会重复处理
2. **错误处理**:处理失败的文件状态会被记录,可以查看错误信息
3. **性能考虑**:大文件会被自动切分,建议根据实际情况调整 chunk_size
4. **扩展性**:可以通过重写 `_load_document_by_type` 方法支持更多文件格式
## 完整示例
参见 `examples/simple_test.py` 获取完整的使用示例。

View File

@ -1,175 +0,0 @@
# BaseRAG 重排功能说明
## 功能概述
BaseRAG现在支持对检索到的文档片段进行重排rerank这是RAG系统中一个重要的优化环节可以提高检索质量和答案的相关性。
## 支持的重排方法
### 1. 相似度重排 (similarity)
- **描述**: 基于余弦相似度的重排方法
- **优点**: 无需额外依赖,使用现有的嵌入模型
- **适用场景**: 轻量级部署,快速原型验证
### 2. CrossEncoder重排 (cross_encoder)
- **描述**: 使用预训练的CrossEncoder模型进行重排
- **依赖**: `sentence-transformers`
- **优点**: 专门训练的重排模型,效果通常更好
- **默认模型**: `cross-encoder/ms-marco-MiniLM-L-6-v2`
### 3. BGE重排 (bge)
- **描述**: 使用BGEBAAI General Embedding重排模型
- **依赖**: `FlagEmbedding`
- **优点**: 中文支持较好,性能优秀
- **默认模型**: `BAAI/bge-reranker-base`
## 支持的嵌入模型类型
### 1. 本地HuggingFace模型 (local)
- **模型名称**: 从HuggingFace Hub下载
- **本地路径**: 使用已下载的本地模型
### 2. 本地部署API接口 (api)
- **描述**: 连接到本地部署的嵌入服务
- **兼容**: OpenAI API格式的本地服务
- **回退机制**: API不可用时自动回退到本地模型
## 使用方法
### 基本配置
```python
from base_rag import BaseRAG
# 嵌入模型配置
embedding_config = {
"type": "local", # 或 "api"
"model_name": "sentence-transformers/all-MiniLM-L6-v2"
}
# 重排配置
rerank_config = {
"enabled": True,
"method": "similarity", # 或 "cross_encoder", "bge"
"top_k": 3 # 重排后返回的文档数量
}
# 初始化RAG系统
rag = MyRAG(
embedding_config=embedding_config,
rerank_config=rerank_config
)
```
### 本地模型配置
```python
# 使用HuggingFace模型名称
embedding_config = {
"type": "local",
"model_name": "sentence-transformers/all-MiniLM-L6-v2"
}
# 使用本地模型路径
embedding_config = {
"type": "local",
"model_path": "/path/to/your/local/model"
}
```
### 本地API接口配置
```python
embedding_config = {
"type": "api",
"api_url": "http://localhost:8000", # 本地API地址
"model": "your-embedding-model",
"api_key": "optional-key" # 可选
}
```
### CrossEncoder重排配置
```python
rerank_config = {
"enabled": True,
"method": "cross_encoder",
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2", # 可选,自定义模型
"top_k": 3
}
```
### BGE重排配置
```python
rerank_config = {
"enabled": True,
"method": "bge",
"model": "BAAI/bge-reranker-base", # 可选,自定义模型
"top_k": 3
}
```
## API方法
### `similarity_search_with_rerank(query, k=None)`
带重排功能的相似性搜索方法,会自动应用配置的重排策略。
```python
# 使用重排搜索
docs = rag.similarity_search_with_rerank("查询问题", k=5)
```
### `similarity_search(query, k=None)`
原始的相似性搜索方法,不使用重排。
```python
# 不使用重排搜索
docs = rag.similarity_search("查询问题", k=5)
```
## 重排工作原理
1. **候选检索**: 首先从向量数据库检索更多的候选文档通常是目标数量的2倍
2. **重排计算**: 根据配置的重排方法,计算查询与每个候选文档的相关性分数
3. **重新排序**: 按照相关性分数对文档进行重新排序
4. **截取结果**: 返回top_k个最相关的文档
## 性能考虑
- **相似度重排**: 计算开销小,速度快
- **CrossEncoder重排**: 需要加载额外模型,计算较慢但效果好
- **BGE重排**: 平衡了效果和性能
- **本地API**: 网络延迟较小,可扩展性好
## 最佳实践
1. **开发阶段**: 使用相似度重排进行快速原型验证
2. **生产环境**: 根据具体需求选择CrossEncoder或BGE重排
3. **分布式部署**: 考虑使用本地API接口
4. **候选文档数量**: 建议设置为最终需要文档数量的2-3倍
5. **模型选择**: 根据语言和领域选择合适的重排模型
## 依赖安装
```bash
# 基本依赖
pip install langchain-huggingface langchain-chroma
# CrossEncoder重排
pip install sentence-transformers
# BGE重排
pip install FlagEmbedding
# 本地API支持可选
pip install langchain-openai
```
## 注意事项
- 重排会增加查询延迟,但通常能显著提高结果质量
- 如果重排模型加载失败,系统会自动回退到相似度重排
- 不同重排方法的效果可能因数据集和查询类型而异
- 本地API接口需要兼容OpenAI API格式
- 建议在实际数据上进行A/B测试来选择最适合的重排方法

Binary file not shown.

View File

@ -1,100 +0,0 @@
#!/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())

View File

@ -83,30 +83,31 @@ async def test_advanced_functionality():
"""测试高级多格式文档和图片功能"""
print("🚀 高级多格式文档和图片内容测试")
print("=" * 60)
# 清理向量数据库
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/advanced_test")
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/storage/chroma_db/ad_test")
if db_path.exists():
shutil.rmtree(db_path)
print("🧹 已清理向量数据库")
# 创建RAG实例 - 启用图片处理
rag = AdvancedTestRAG(
vector_store_name="advanced_test",
persist_directory="./storage/chroma_db/ad_test",
retriever_top_k=5,
storage_directory="/Users/liruwei/Documents/code/project/demo/base_rag/test_files",
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/advanced_test_status.db",
storage_directory="/Users/liruwei/Documents/code/project/demo/base_rag/storage/files",
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/storage/status_db/advanced_test_status.db",
# 启用图片处理 - 使用本地BLIP模型获得更好的图片文本识别
image_config={
"enabled": True,
"type": "local",
"model": "Salesforce/blip-image-captioning-base"
}
"model": "Salesforce/blip-image-captioning-base",
},
)
print("✅ 高级RAG实例创建成功 (已启用图片处理)")
print()
# 测试多格式文档
test_files = [
{
@ -140,77 +141,77 @@ async def test_advanced_functionality():
"expect_images": False
}
]
# 筛选存在的文件
test_dir = Path("/Users/liruwei/Documents/code/project/demo/base_rag/test_files")
available_files = []
for file_info in test_files:
if (test_dir / file_info["file"]).exists():
available_files.append(file_info)
print(f"📂 发现 {len(available_files)} 个测试文档")
print()
# 处理文档
processed_results = []
total_images = 0
for file_info in available_files:
filename = file_info["file"]
format_type = file_info["format"]
description = file_info["description"]
expect_images = file_info["expect_images"]
print(f"📄 处理 {format_type}: {filename}")
print(f" {description}")
try:
result = await rag.ingest(str(test_dir / filename))
if result and result.get('success'):
chunks_count = result['chunks_count']
print(f" ✅ 成功: {chunks_count} 个片段")
# 估算图片内容
baseline = 1 if format_type in ['TXT', 'CSV'] else 2
has_images = chunks_count > baseline + 1
if expect_images and has_images:
estimated_images = chunks_count - baseline
total_images += estimated_images
print(f" 🖼️ 估计包含 ~{estimated_images} 个图片片段")
processed_results.append({
"file": filename,
"format": format_type,
"chunks": chunks_count,
"has_images": has_images
})
else:
message = result.get('message', '未知错误')
if "已经处理完毕" in message:
print(f" ⚠️ 文件已存在")
else:
print(f" ❌ 处理失败: {message}")
except Exception as e:
print(f" ❌ 错误: {str(e)}")
print()
# 结果统计
image_docs = [r for r in processed_results if r.get("has_images")]
text_docs = [r for r in processed_results if not r.get("has_images")]
print("📊 处理结果统计:")
print(f" 📄 纯文本文档: {len(text_docs)}")
print(f" 🖼️ 含图片文档: {len(image_docs)}")
if total_images > 0:
print(f" 📸 估计图片总数: ~{total_images}")
print()
# 高级查询测试
print("🔍 高级查询测试...")
test_queries = [
{
"question": "数据科学的核心技术有哪些?",
@ -237,16 +238,16 @@ async def test_advanced_functionality():
"focus": "研究内容"
}
]
image_content_found = False
for i, query_info in enumerate(test_queries, 1):
question = query_info["question"]
focus = query_info["focus"]
print(f"\n❓ 查询 {i}: {question}")
print(f" 🎯 重点: {focus}")
try:
answer = await rag.query(question)
if "抱歉" not in answer:
@ -254,22 +255,22 @@ async def test_advanced_functionality():
if "🖼️ [图片" in answer:
print(f" 🖼️ ✅ 检索到图片内容!")
image_content_found = True
# 分析结果
lines = answer.split('\n')
if lines:
source_line = lines[0] if lines[0].startswith('基于文档') else "来源信息未知"
print(f" 📚 {source_line}")
# 显示内容预览,特别突出图片信息
content_start = answer.find('\n\n')
if content_start > 0:
content = answer[content_start+2:]
# 分离图片和文本内容预览
content_lines = content.split('\n\n')
preview_parts = []
for line in content_lines[:2]: # 只显示前2个部分
if "🖼️ [图片" in line:
# 图片内容特殊处理
@ -279,17 +280,17 @@ async def test_advanced_functionality():
# 文本内容
text_preview = line[:100] + "..." if len(line) > 100 else line
preview_parts.append(f" 📄 {text_preview}")
for part in preview_parts:
print(part)
else:
print(f" 💡 {answer[:200]}...")
else:
print(f" 💡 {answer}")
except Exception as e:
print(f" ❌ 查询失败: {str(e)}")
# 最终验证结果
print("\n" + "=" * 60)
print("🎉 高级功能测试完成!")

View File

@ -54,40 +54,40 @@ async def test_basic_functionality():
"""测试基础RAG功能"""
print("🔧 基础RAG功能测试")
print("=" * 50)
# 清理向量数据库
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/simple_test")
if db_path.exists():
shutil.rmtree(db_path)
print("🧹 已清理向量数据库")
# 创建RAG实例 - 禁用图片处理用于基础测试
rag = SimpleTestRAG(
vector_store_name="simple_test",
retriever_top_k=3,
storage_directory="/Users/liruwei/Documents/code/project/demo/base_rag/test_files",
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/simple_test_status.db",
image_config={"enabled": False} # 基础测试禁用图片
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/status_db/simple_test_status.db",
image_config={"enabled": False}, # 基础测试禁用图片
)
print("✅ RAG实例创建成功")
print()
# 测试基础文档
test_files = ["test_document.txt", "test_markdown.md", "python_basics.txt", "data_science.txt"]
print("📂 处理基础文档...")
processed_count = 0
for filename in test_files:
file_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/test_files") / filename
if not file_path.exists():
print(f"⚠️ {filename} - 文件不存在,跳过")
continue
print(f"📄 处理: {filename}")
try:
result = await rag.ingest(str(file_path))
if result and result.get('success'):
@ -102,23 +102,23 @@ async def test_basic_functionality():
print(f" ❌ 失败: {message}")
except Exception as e:
print(f" ❌ 错误: {str(e)}")
print(f"\n📊 处理完成: {processed_count}/{len(test_files)} 个文件")
print()
# 基础查询测试
print("🔍 基础查询测试...")
test_queries = [
"Python编程语言的特点",
"数据科学的核心技术",
"机器学习的应用",
"什么是深度学习"
]
for i, question in enumerate(test_queries, 1):
print(f"\n❓ 查询 {i}: {question}")
try:
answer = await rag.query(question)
if "抱歉" not in answer:
@ -126,7 +126,7 @@ async def test_basic_functionality():
lines = answer.split('\n')
source_line = lines[0] if lines[0].startswith('基于文档') else "来源未知"
print(f" 📚 {source_line}")
# 显示内容预览
content_start = answer.find('\n\n')
if content_start > 0:
@ -139,7 +139,7 @@ async def test_basic_functionality():
print(f" 💡 {answer}")
except Exception as e:
print(f" ❌ 查询失败: {str(e)}")
print("\n" + "=" * 50)
print("🎉 基础功能测试完成!")
print("✅ 验证项目:")

BIN
status.db

Binary file not shown.

File diff suppressed because one or more lines are too long

Binary file not shown.

View File

@ -1,6 +0,0 @@
日期,产品,销售额,数量,客户类型,销售员
2024-01-01,笔记本电脑,8500,5,企业,张三
2024-01-02,台式机,6200,4,个人,李四
2024-01-03,平板电脑,3200,8,学生,王五
2024-01-04,智能手机,4500,9,个人,张三
2024-01-05,耳机,280,12,学生,李四
1 日期 产品 销售额 数量 客户类型 销售员
2 2024-01-01 笔记本电脑 8500 5 企业 张三
3 2024-01-02 台式机 6200 4 个人 李四
4 2024-01-03 平板电脑 3200 8 学生 王五
5 2024-01-04 智能手机 4500 9 个人 张三
6 2024-01-05 耳机 280 12 学生 李四

View File

@ -1,106 +0,0 @@
#!/usr/bin/env python3
"""
测试优化后的图片OCR功能
"""
import sys
import os
import asyncio
from pathlib import Path
# 添加源码路径
sys.path.append(os.path.join(os.path.dirname(__file__), "src"))
from base_rag.image_processor import ImageProcessor
async def test_image_ocr():
"""测试不同模式下的图片OCR功能"""
# 测试配置
configs = [
{"type": "local", "engine": "easyocr"},
{"type": "basic"},
# {"type": "api", "api_url": "http://localhost:8000/image2text"} # 需要实际API
]
print("🧪 开始测试图片OCR功能")
print("=" * 50)
# 寻找测试图片
test_images = []
# 检查常见的图片位置
possible_paths = [
"./test_files/",
"./demo_documents/",
"./examples/",
"./"
]
image_extensions = ['.png', '.jpg', '.jpeg', '.gif', '.bmp']
for path in possible_paths:
if os.path.exists(path):
for file in os.listdir(path):
if any(file.lower().endswith(ext) for ext in image_extensions):
test_images.append(os.path.join(path, file))
if not test_images:
print("⚠️ 未找到测试图片,创建示例图片...")
# 创建一个简单的测试图片
try:
from PIL import Image, ImageDraw, ImageFont
# 创建包含文字的测试图片
img = Image.new('RGB', (400, 200), color='white')
draw = ImageDraw.Draw(img)
# 尝试使用默认字体
try:
font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", 24)
except:
font = ImageFont.load_default()
# 添加测试文字
test_text = "Hello World!\nPython OCR Test\n测试中文识别"
draw.text((50, 50), test_text, fill='black', font=font)
test_image_path = "./test_ocr_image.png"
img.save(test_image_path)
test_images = [test_image_path]
print(f"✅ 创建测试图片: {test_image_path}")
except Exception as e:
print(f"❌ 创建测试图片失败: {e}")
return
print(f"📸 找到 {len(test_images)} 个测试图片")
# 测试每种配置
for i, config in enumerate(configs, 1):
print(f"\n🔧 测试配置 {i}: {config}")
print("-" * 30)
try:
processor = ImageProcessor(config)
# 处理每个测试图片
for img_path in test_images[:2]: # 限制测试图片数量
print(f"\n📷 处理图片: {os.path.basename(img_path)}")
if os.path.exists(img_path):
result = processor.extract_image_description(img_path)
print(f"结果:\n{result}")
else:
print(f"❌ 图片不存在: {img_path}")
except Exception as e:
print(f"❌ 配置 {config} 测试失败: {e}")
print("\n" + "=" * 50)
print("🏁 测试完成")
if __name__ == "__main__":
asyncio.run(test_image_ocr())