Compare commits
2 Commits
c6e020a170
...
1e2284728f
Author | SHA1 | Date |
---|---|---|
|
1e2284728f | |
|
5d50b9a23e |
|
@ -0,0 +1,88 @@
|
|||
# Examples 示例文件
|
||||
|
||||
本目录包含两个主要的测试示例:
|
||||
|
||||
## 📁 simple_test.py - 基础功能测试
|
||||
**用途**: 验证RAG系统的基础功能
|
||||
- 🔧 纯文本文档处理
|
||||
- 📄 文档加载和切分
|
||||
- 🔍 文本向量化和存储
|
||||
- 🔎 相似性搜索
|
||||
- 📝 查询结果整合
|
||||
|
||||
**特点**:
|
||||
- 禁用图片处理(专注基础功能)
|
||||
- 适合快速验证系统可用性
|
||||
- 轻量级测试
|
||||
|
||||
**运行**:
|
||||
```bash
|
||||
python examples/simple_test.py
|
||||
```
|
||||
|
||||
## 🚀 ad_test.py - 高级功能测试
|
||||
**用途**: 验证多格式文档和图片内容识别
|
||||
- 📄 多格式文档解析 (DOCX, PDF, XLSX, CSV)
|
||||
- 🖼️ 图片自动提取和处理
|
||||
- 🤖 图片内容描述生成
|
||||
- 📝 图片文本内容识别 (OCR)
|
||||
- 🔍 混合内容检索 (文本+图片)
|
||||
- 📊 内容分类显示
|
||||
|
||||
**特点**:
|
||||
- 启用完整图片处理功能
|
||||
- 使用BLIP模型进行图片理解
|
||||
- 支持图片中文本提取
|
||||
- 增强的查询结果显示
|
||||
|
||||
**运行**:
|
||||
```bash
|
||||
python examples/ad_test.py
|
||||
```
|
||||
|
||||
## 🔧 图片文本识别功能
|
||||
|
||||
高级测试(`ad_test.py`)包含增强的图片文本识别功能:
|
||||
|
||||
### ✅ 图片内容处理
|
||||
- **自动提取**: 从DOCX和PDF文档中自动提取嵌入的图片
|
||||
- **智能描述**: 使用BLIP模型生成图片内容描述
|
||||
- **文本识别**: 支持OCR提取图片中的文字内容
|
||||
- **分类标记**: 自动识别图片类型(技术图、数据图表等)
|
||||
|
||||
### 📝 OCR文本提取
|
||||
系统尝试从图片中提取文字内容,支持:
|
||||
- **pytesseract**: 高精度OCR引擎(需要安装)
|
||||
- **easyocr**: 备用OCR方案(支持中英文)
|
||||
- **基础模式**: 如果OCR库不可用,提供基础信息
|
||||
|
||||
### 🔍 增强检索体验
|
||||
- **内容分类**: 查询结果区分图片内容和文本内容
|
||||
- **统计信息**: 显示检索到的文本和图片数量
|
||||
- **格式化显示**: 图片内容带特殊标记 `🖼️ [图片内容]`
|
||||
|
||||
## 📋 测试文档要求
|
||||
|
||||
### 基础测试文档
|
||||
- `python_basics.txt` - Python基础知识
|
||||
- `data_science.txt` - 数据科学内容
|
||||
|
||||
### 高级测试文档
|
||||
- `complex_data_science.docx` - 包含图片的Word文档
|
||||
- `ai_research_report.pdf` - 包含图片的PDF报告
|
||||
- `company_report.xlsx` - Excel工作簿
|
||||
- `sales_data.csv` - CSV数据文件
|
||||
|
||||
## 🎯 预期效果
|
||||
|
||||
### 基础测试
|
||||
- ✅ 文档正常加载和处理
|
||||
- ✅ 文本查询返回相关结果
|
||||
- ✅ 系统响应时间正常
|
||||
|
||||
### 高级测试
|
||||
- ✅ 多格式文档成功解析
|
||||
- ✅ 图片内容被自动识别和描述
|
||||
- 🖼️ 图片查询能返回图片相关内容
|
||||
- 📊 查询结果包含内容类型统计
|
||||
- 🔍 图片和文本内容可被统一检索
|
|
@ -0,0 +1,320 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
高级测试示例 - 多格式文档和图片内容识别
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
|
||||
# 过滤掉PyTorch的FutureWarning
|
||||
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
|
||||
|
||||
# 添加源码路径
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
|
||||
|
||||
from base_rag.core import BaseRAG
|
||||
|
||||
|
||||
class AdvancedTestRAG(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=5)
|
||||
|
||||
if not docs:
|
||||
return "抱歉,没有找到相关信息。"
|
||||
|
||||
# 分析和整理搜索结果
|
||||
sources = []
|
||||
contexts = []
|
||||
image_count = 0
|
||||
text_count = 0
|
||||
|
||||
for doc in docs:
|
||||
source = doc.metadata.get("source_file", "未知来源")
|
||||
doc_type = doc.metadata.get("type", "text")
|
||||
content = doc.page_content.strip()
|
||||
|
||||
if source not in sources:
|
||||
sources.append(source)
|
||||
|
||||
# 处理不同类型的内容
|
||||
if doc_type == "image":
|
||||
# 增强图片内容显示
|
||||
image_count += 1
|
||||
enhanced_content = f"🖼️ [图片 {image_count}] {content}"
|
||||
|
||||
# 如果图片描述中包含文件信息,提取并格式化
|
||||
if "图片文件:" in content and "尺寸:" in content:
|
||||
parts = content.split(" | ")
|
||||
if len(parts) >= 3:
|
||||
file_info = parts[0].replace("图片文件: ", "")
|
||||
size_info = parts[1].replace("尺寸: ", "")
|
||||
type_info = parts[2].replace("类型: ", "")
|
||||
enhanced_content = f"🖼️ [图片内容] {file_info}\n 📐 尺寸: {size_info} | 🏷️ 类型: {type_info}"
|
||||
|
||||
contexts.append(enhanced_content)
|
||||
else:
|
||||
text_count += 1
|
||||
contexts.append(f"📄 {content}")
|
||||
|
||||
context = "\n\n".join(contexts)
|
||||
sources_str = "、".join(sources)
|
||||
|
||||
# 添加内容统计信息
|
||||
stats = f"({text_count}文本"
|
||||
if image_count > 0:
|
||||
stats += f" + {image_count}图片"
|
||||
stats += ")"
|
||||
|
||||
return f"基于文档({sources_str}){stats}的信息:\n\n{context}"
|
||||
|
||||
|
||||
async def test_advanced_functionality():
|
||||
"""测试高级多格式文档和图片功能"""
|
||||
print("🚀 高级多格式文档和图片内容测试")
|
||||
print("=" * 60)
|
||||
|
||||
# 清理向量数据库
|
||||
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/advanced_test")
|
||||
if db_path.exists():
|
||||
shutil.rmtree(db_path)
|
||||
print("🧹 已清理向量数据库")
|
||||
|
||||
# 创建RAG实例 - 启用图片处理
|
||||
rag = AdvancedTestRAG(
|
||||
vector_store_name="advanced_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",
|
||||
# 启用图片处理 - 使用本地BLIP模型获得更好的图片文本识别
|
||||
image_config={
|
||||
"enabled": True,
|
||||
"type": "local",
|
||||
"model": "Salesforce/blip-image-captioning-base"
|
||||
}
|
||||
)
|
||||
|
||||
print("✅ 高级RAG实例创建成功 (已启用图片处理)")
|
||||
print()
|
||||
|
||||
# 测试多格式文档
|
||||
test_files = [
|
||||
{
|
||||
"file": "test_document.txt",
|
||||
"format": "TXT",
|
||||
"description": "纯文本文档",
|
||||
"expect_images": False
|
||||
},
|
||||
{
|
||||
"file": "complex_data_science.docx",
|
||||
"format": "DOCX",
|
||||
"description": "Word文档(含图片)",
|
||||
"expect_images": True
|
||||
},
|
||||
{
|
||||
"file": "ai_research_report.pdf",
|
||||
"format": "PDF",
|
||||
"description": "PDF报告(含图片)",
|
||||
"expect_images": True
|
||||
},
|
||||
{
|
||||
"file": "company_report.xlsx",
|
||||
"format": "XLSX",
|
||||
"description": "Excel工作簿",
|
||||
"expect_images": False
|
||||
},
|
||||
{
|
||||
"file": "sales_data.csv",
|
||||
"format": "CSV",
|
||||
"description": "CSV数据文件",
|
||||
"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": "数据科学的核心技术有哪些?",
|
||||
"focus": "文本内容"
|
||||
},
|
||||
{
|
||||
"question": "文档中的图片显示了什么内容?",
|
||||
"focus": "图片内容"
|
||||
},
|
||||
{
|
||||
"question": "Python生态系统相关的信息",
|
||||
"focus": "综合内容"
|
||||
},
|
||||
{
|
||||
"question": "销售数据分析结果",
|
||||
"focus": "数据内容"
|
||||
},
|
||||
{
|
||||
"question": "技术架构或框架图的内容",
|
||||
"focus": "图片技术内容"
|
||||
},
|
||||
{
|
||||
"question": "人工智能研究的挑战和机遇",
|
||||
"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:
|
||||
# 检查是否包含图片内容
|
||||
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:
|
||||
# 图片内容特殊处理
|
||||
img_preview = line[:200] + "..." if len(line) > 200 else line
|
||||
preview_parts.append(f" 🖼️ {img_preview}")
|
||||
else:
|
||||
# 文本内容
|
||||
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("🎉 高级功能测试完成!")
|
||||
print()
|
||||
print("✅ 功能验证结果:")
|
||||
print(" 📄 多格式文档解析 - ✅")
|
||||
print(" 🖼️ 图片自动提取 - ✅" if image_docs else " 🖼️ 图片自动提取 - ⚠️")
|
||||
print(" 🤖 图片文本识别 - ✅" if image_content_found else " 🤖 图片文本识别 - ⚠️")
|
||||
print(" 🔍 混合内容检索 - ✅" if image_content_found else " 🔍 混合内容检索 - ⚠️")
|
||||
print(" 📊 内容分类显示 - ✅")
|
||||
print()
|
||||
print("🔧 支持的格式:")
|
||||
for file_info in available_files:
|
||||
icon = "🖼️" if file_info["expect_images"] else "📄"
|
||||
print(f" {icon} {file_info['format']} - {file_info['description']}")
|
||||
print()
|
||||
print("💡 图片文本识别特性:")
|
||||
if image_content_found:
|
||||
print(" ✅ 自动提取图片中的视觉信息")
|
||||
print(" ✅ 生成图片内容描述文本")
|
||||
print(" ✅ 图片信息可被向量化和检索")
|
||||
print(" ✅ 支持图片尺寸和类型识别")
|
||||
else:
|
||||
print(" ⚠️ 需要包含图片的测试文档验证")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(test_advanced_functionality())
|
|
@ -1,207 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
完整的多格式文件测试 - 包含图片的 DOCX、PDF、Excel、CSV
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
|
||||
# 过滤掉PyTorch的FutureWarning
|
||||
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
|
||||
|
||||
# 添加源码路径
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
|
||||
|
||||
from base_rag.core import BaseRAG
|
||||
|
||||
|
||||
class AdvancedFormatRAG(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 "抱歉,没有找到相关信息。"
|
||||
|
||||
# 显示搜索到的文档来源
|
||||
sources = []
|
||||
contexts = []
|
||||
for doc in docs:
|
||||
source = doc.metadata.get("source_file", "未知来源")
|
||||
content = doc.page_content.strip()
|
||||
|
||||
if source not in sources:
|
||||
sources.append(source)
|
||||
contexts.append(content)
|
||||
|
||||
context = "\n\n".join(contexts)
|
||||
sources_str = "、".join(sources)
|
||||
|
||||
return f"基于以下文档({sources_str})的信息:\n\n{context}"
|
||||
|
||||
|
||||
async def test_advanced_formats():
|
||||
"""测试高级文件格式处理"""
|
||||
print("🚀 高级多格式文件处理测试")
|
||||
print("=" * 60)
|
||||
|
||||
# 清理旧的向量数据库
|
||||
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/advanced_formats")
|
||||
if db_path.exists():
|
||||
shutil.rmtree(db_path)
|
||||
print("🧹 已清理旧的向量数据库")
|
||||
|
||||
# 创建RAG实例
|
||||
rag = AdvancedFormatRAG(
|
||||
vector_store_name="advanced_formats",
|
||||
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/advanced_status.db",
|
||||
)
|
||||
|
||||
# 测试文件列表 - 包含新创建的文件
|
||||
test_files = [
|
||||
{
|
||||
"file": "complex_data_science.docx",
|
||||
"format": "DOCX",
|
||||
"description": "复杂Word文档(含表格和图片)"
|
||||
},
|
||||
{
|
||||
"file": "sales_data.csv",
|
||||
"format": "CSV",
|
||||
"description": "销售数据CSV文件"
|
||||
},
|
||||
{
|
||||
"file": "company_report.xlsx",
|
||||
"format": "XLSX",
|
||||
"description": "多工作表Excel文件"
|
||||
},
|
||||
{
|
||||
"file": "ai_research_report.pdf",
|
||||
"format": "PDF",
|
||||
"description": "AI研究报告PDF(含图片)"
|
||||
}
|
||||
]
|
||||
|
||||
print("📂 处理高级格式文件...")
|
||||
processed_count = 0
|
||||
|
||||
for file_info in test_files:
|
||||
filename = file_info["file"]
|
||||
format_type = file_info["format"]
|
||||
description = file_info["description"]
|
||||
|
||||
file_path = Path("../test_files") / filename
|
||||
|
||||
if not file_path.exists():
|
||||
# 尝试绝对路径
|
||||
file_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/test_files") / filename
|
||||
|
||||
if not file_path.exists():
|
||||
print(f"❌ {format_type}: {filename} - 文件不存在")
|
||||
continue
|
||||
|
||||
print(f"📄 处理 {format_type}: {filename}")
|
||||
print(f" {description}")
|
||||
|
||||
try:
|
||||
result = await rag.ingest(str(file_path))
|
||||
if result and result.get('success'):
|
||||
print(f" ✅ 成功: {result['chunks_count']} 个片段")
|
||||
processed_count += 1
|
||||
else:
|
||||
print(f" ⚠️ 跳过: {result.get('message', '可能已存在')}")
|
||||
if "已经处理完毕" in str(result.get('message', '')):
|
||||
processed_count += 1
|
||||
except Exception as e:
|
||||
print(f" ❌ 失败: {str(e)}")
|
||||
print()
|
||||
|
||||
print(f"📊 处理完成: {processed_count}/{len(test_files)} 个文件")
|
||||
print()
|
||||
|
||||
# 测试针对性查询
|
||||
print("💬 高级格式查询测试...")
|
||||
|
||||
queries = [
|
||||
{
|
||||
"question": "数据科学的核心技术有哪些?",
|
||||
"expected": "complex_data_science.docx"
|
||||
},
|
||||
{
|
||||
"question": "销售数据中哪个产品销售额最高?",
|
||||
"expected": "sales_data.csv"
|
||||
},
|
||||
{
|
||||
"question": "公司员工信息包含哪些部门?",
|
||||
"expected": "company_report.xlsx"
|
||||
},
|
||||
{
|
||||
"question": "人工智能研究面临的挑战是什么?",
|
||||
"expected": "ai_research_report.pdf"
|
||||
},
|
||||
{
|
||||
"question": "Python在数据科学中的作用?",
|
||||
"expected": "多个文档"
|
||||
}
|
||||
]
|
||||
|
||||
for i, query_info in enumerate(queries, 1):
|
||||
question = query_info["question"]
|
||||
expected = query_info["expected"]
|
||||
|
||||
print(f"\n❓ 查询 {i}: {question}")
|
||||
print(f" 期望来源: {expected}")
|
||||
|
||||
try:
|
||||
answer = await rag.query(question)
|
||||
if "抱歉" not in answer:
|
||||
# 分离来源信息和内容
|
||||
parts = answer.split('\n\n', 1)
|
||||
if len(parts) == 2:
|
||||
source_info = parts[0]
|
||||
content = parts[1]
|
||||
|
||||
print(f" 📚 {source_info}")
|
||||
|
||||
# 显示内容摘要(前150字符)
|
||||
if len(content) > 150:
|
||||
content_preview = content[:150] + "..."
|
||||
else:
|
||||
content_preview = content
|
||||
|
||||
print(f" 💡 {content_preview}")
|
||||
else:
|
||||
print(f" 💡 {answer[:150]}...")
|
||||
else:
|
||||
print(f" 💡 {answer}")
|
||||
except Exception as e:
|
||||
print(f" ❌ 查询失败: {str(e)}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("🎉 高级多格式文件测试完成!")
|
||||
print("✅ 支持的格式:")
|
||||
print(" 📄 DOCX - Word文档 (含表格、图片)")
|
||||
print(" 📊 CSV - 逗号分隔值文件")
|
||||
print(" 📈 XLSX - Excel工作簿 (多工作表)")
|
||||
print(" 📑 PDF - 便携式文档格式 (含图片)")
|
||||
print()
|
||||
print("🔧 技术特性:")
|
||||
print(" 🔄 异步处理 - 非阻塞I/O操作")
|
||||
print(" 🧠 智能解析 - 自动识别文件格式")
|
||||
print(" 🔍 跨格式查询 - 统一检索接口")
|
||||
print(" 📋 表格数据提取 - 结构化信息处理")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(test_advanced_formats())
|
|
@ -1,186 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
多格式文件测试 - 测试 TXT、MD、DOCX 文件格式
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
# 过滤掉PyTorch的FutureWarning
|
||||
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
|
||||
|
||||
# 添加源码路径
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
|
||||
|
||||
from base_rag.core import BaseRAG
|
||||
|
||||
|
||||
class MultiFormatRAG(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 "抱歉,没有找到相关信息。"
|
||||
|
||||
# 显示搜索到的文档来源
|
||||
sources = []
|
||||
contexts = []
|
||||
for doc in docs:
|
||||
source = doc.metadata.get("source_file", "未知来源")
|
||||
content = doc.page_content.strip()
|
||||
|
||||
if source not in sources:
|
||||
sources.append(source)
|
||||
contexts.append(content)
|
||||
|
||||
context = "\n\n".join(contexts)
|
||||
sources_str = "、".join(sources)
|
||||
|
||||
return f"基于以下文档({sources_str})的信息:\n\n{context}"
|
||||
|
||||
|
||||
async def test_multiple_formats():
|
||||
"""测试多种文件格式处理"""
|
||||
print("🚀 多格式文件处理测试")
|
||||
print("=" * 50)
|
||||
|
||||
# 创建RAG实例
|
||||
rag = MultiFormatRAG(
|
||||
vector_store_name="multiformat_kb",
|
||||
retriever_top_k=3,
|
||||
storage_directory="../test_files", # 相对于examples目录
|
||||
status_db_path="../status.db", # 相对于examples目录
|
||||
)
|
||||
|
||||
# 测试文件列表
|
||||
test_files = [
|
||||
{
|
||||
"file": "knowledge.txt",
|
||||
"format": "TXT",
|
||||
"description": "纯文本文件"
|
||||
},
|
||||
{
|
||||
"file": "python_guide.md",
|
||||
"format": "MD",
|
||||
"description": "Markdown文件"
|
||||
},
|
||||
{
|
||||
"file": "machine_learning.md",
|
||||
"format": "MD",
|
||||
"description": "Markdown文件"
|
||||
},
|
||||
{
|
||||
"file": "deep_learning_guide.docx",
|
||||
"format": "DOCX",
|
||||
"description": "Word文档"
|
||||
},
|
||||
{
|
||||
"file": "complex_data_science.docx",
|
||||
"format": "DOCX",
|
||||
"description": "复杂Word文档(含表格)"
|
||||
},
|
||||
{
|
||||
"file": "sales_data.csv",
|
||||
"format": "CSV",
|
||||
"description": "CSV数据文件"
|
||||
},
|
||||
{
|
||||
"file": "company_report.xlsx",
|
||||
"format": "XLSX",
|
||||
"description": "Excel工作簿"
|
||||
},
|
||||
{
|
||||
"file": "ai_research_report.pdf",
|
||||
"format": "PDF",
|
||||
"description": "PDF文档"
|
||||
}
|
||||
]
|
||||
|
||||
print("📂 处理文件...")
|
||||
processed_count = 0
|
||||
|
||||
for file_info in test_files:
|
||||
filename = file_info["file"]
|
||||
format_type = file_info["format"]
|
||||
description = file_info["description"]
|
||||
|
||||
file_path = Path("../test_files") / filename
|
||||
|
||||
if not file_path.exists():
|
||||
print(f"❌ {format_type}: {filename} - 文件不存在")
|
||||
continue
|
||||
|
||||
print(f"📄 处理 {format_type}: {filename} ({description})")
|
||||
|
||||
try:
|
||||
result = await rag.ingest(str(file_path))
|
||||
if result and result.get('success'):
|
||||
print(f" ✅ 成功: {result['chunks_count']} 个片段")
|
||||
processed_count += 1
|
||||
else:
|
||||
print(f" ⚠️ 跳过: {result.get('message', '可能已存在')}")
|
||||
processed_count += 1 # 已存在也算处理过
|
||||
except Exception as e:
|
||||
print(f" ❌ 失败: {str(e)}")
|
||||
|
||||
print(f"\n📊 处理完成: {processed_count}/{len(test_files)} 个文件")
|
||||
print()
|
||||
|
||||
# 测试跨格式查询
|
||||
print("💬 跨格式查询测试...")
|
||||
|
||||
queries = [
|
||||
"Python有什么特点?",
|
||||
"什么是机器学习?",
|
||||
"深度学习的应用领域有哪些?",
|
||||
"数据科学的核心技术有哪些?",
|
||||
"销售数据中哪个产品销售额最高?",
|
||||
"公司员工的平均年薪是多少?",
|
||||
"人工智能的主要挑战是什么?",
|
||||
"机器学习有哪些类型?"
|
||||
]
|
||||
|
||||
for query in queries:
|
||||
print(f"\n❓ {query}")
|
||||
try:
|
||||
answer = await rag.query(query)
|
||||
if "抱歉" not in answer:
|
||||
# 分离来源信息和内容
|
||||
parts = answer.split('\n\n', 1)
|
||||
if len(parts) == 2:
|
||||
source_info = parts[0] # "基于以下文档..."
|
||||
content = parts[1] # 实际内容
|
||||
|
||||
print(f" 📚 {source_info}")
|
||||
|
||||
# 显示内容摘要(前200字符)
|
||||
if len(content) > 200:
|
||||
content_preview = content[:200] + "..."
|
||||
else:
|
||||
content_preview = content
|
||||
|
||||
print(f" 💡 {content_preview}")
|
||||
else:
|
||||
print(f" 💡 {answer}")
|
||||
else:
|
||||
print(f" 💡 {answer}")
|
||||
except Exception as e:
|
||||
print(f" ❌ 查询失败: {str(e)}")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("✅ 多格式文件测试完成!")
|
||||
print("支持的格式: TXT, MD, DOCX, CSV, XLSX, PDF")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(test_multiple_formats())
|
|
@ -1,6 +1,6 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
简单的文件处理测试
|
||||
简单测试示例 - 基础RAG功能验证
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
@ -8,6 +8,7 @@ import os
|
|||
import asyncio
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
|
||||
# 过滤掉PyTorch的FutureWarning
|
||||
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
|
||||
|
@ -15,90 +16,138 @@ warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
|
|||
# 添加源码路径
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
|
||||
|
||||
from base_rag.core import BaseRAG, FileStatus
|
||||
from base_rag.core import BaseRAG
|
||||
|
||||
|
||||
class SimpleRAG(BaseRAG):
|
||||
"""简单的RAG实现示例"""
|
||||
class SimpleTestRAG(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=2)
|
||||
"""查询实现"""
|
||||
docs = await self.similarity_search_with_rerank(question, k=3)
|
||||
|
||||
if not docs:
|
||||
return "抱歉,没有找到相关信息。"
|
||||
|
||||
# 显示搜索到的文档来源
|
||||
# 整理搜索结果
|
||||
sources = []
|
||||
contexts = []
|
||||
for doc in docs:
|
||||
source = doc.metadata.get("source_file", "未知来源")
|
||||
content = doc.page_content.strip()
|
||||
|
||||
if source not in sources:
|
||||
sources.append(source)
|
||||
contexts.append(doc.page_content.strip())
|
||||
contexts.append(content)
|
||||
|
||||
context = "\n\n".join(contexts)
|
||||
sources_str = "、".join(sources)
|
||||
|
||||
return f"基于以下文档({sources_str})的信息:\n\n{context}"
|
||||
return f"基于文档({sources_str})的信息:\n\n{context}"
|
||||
|
||||
|
||||
async def test_file_processing():
|
||||
print("=== 文件处理功能测试 ===\n")
|
||||
|
||||
# 创建RAG实例
|
||||
rag = SimpleRAG(
|
||||
vector_store_name="test_kb",
|
||||
retriever_top_k=2,
|
||||
storage_directory="./test_files", # 统一使用test_files目录
|
||||
status_db_path="./status.db", # 统一数据库名称
|
||||
)
|
||||
|
||||
# 使用现有的测试文件
|
||||
test_dir = Path("./test_files")
|
||||
async def test_basic_functionality():
|
||||
"""测试基础RAG功能"""
|
||||
print("🔧 基础RAG功能测试")
|
||||
print("=" * 50)
|
||||
|
||||
# 使用已有的测试文件
|
||||
python_file = test_dir / "python_basics.txt"
|
||||
web_file = test_dir / "web_frameworks.txt"
|
||||
datascience_file = test_dir / "data_science.txt"
|
||||
|
||||
print("1. 处理多个知识文件...")
|
||||
files_to_process = [python_file, web_file, datascience_file]
|
||||
|
||||
for file_path in files_to_process:
|
||||
result = await rag.ingest(str(file_path), chunk_size=200, chunk_overlap=20)
|
||||
print(
|
||||
f"处理 {file_path.name}: {result['message']} (片段数: {result.get('chunks_count', 0)})"
|
||||
)
|
||||
# 清理向量数据库
|
||||
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} # 基础测试禁用图片
|
||||
)
|
||||
|
||||
print("✅ RAG实例创建成功")
|
||||
print()
|
||||
|
||||
print("2. 查询测试...")
|
||||
questions = [
|
||||
"Python是谁创建的?",
|
||||
"Flask和Django有什么区别?",
|
||||
"Pandas是做什么的?",
|
||||
"什么是NumPy?",
|
||||
"FastAPI有什么特点?",
|
||||
|
||||
# 测试基础文档
|
||||
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'):
|
||||
print(f" ✅ 成功: {result['chunks_count']} 个片段")
|
||||
processed_count += 1
|
||||
else:
|
||||
message = result.get('message', '未知错误')
|
||||
if "已经处理完毕" in message:
|
||||
print(f" ⚠️ 已存在,跳过")
|
||||
processed_count += 1
|
||||
else:
|
||||
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 question in questions:
|
||||
print(f"问题: {question}")
|
||||
answer = await rag.query(question)
|
||||
print(f"回答: {answer[:150]}...")
|
||||
print("-" * 50)
|
||||
print()
|
||||
|
||||
print("3. 查看文件状态...")
|
||||
files = await rag.get_file_processing_status()
|
||||
for file_info in files:
|
||||
print(f"文件: {file_info['filename']} | 状态: {file_info['status']}")
|
||||
|
||||
print("\n=== 测试完成 ===")
|
||||
|
||||
for i, question in enumerate(test_queries, 1):
|
||||
print(f"\n❓ 查询 {i}: {question}")
|
||||
|
||||
try:
|
||||
answer = await rag.query(question)
|
||||
if "抱歉" not in answer:
|
||||
# 显示结果摘要
|
||||
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:
|
||||
content = answer[content_start+2:]
|
||||
preview = content[:150] + "..." if len(content) > 150 else content
|
||||
print(f" 💡 {preview}")
|
||||
else:
|
||||
print(f" 💡 {answer[:150]}...")
|
||||
else:
|
||||
print(f" 💡 {answer}")
|
||||
except Exception as e:
|
||||
print(f" ❌ 查询失败: {str(e)}")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("🎉 基础功能测试完成!")
|
||||
print("✅ 验证项目:")
|
||||
print(" 📄 文档加载和切分")
|
||||
print(" 🔍 文本向量化和存储")
|
||||
print(" 🔎 相似性搜索")
|
||||
print(" 📝 查询结果整合")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(test_file_processing())
|
||||
asyncio.run(test_basic_functionality())
|
||||
|
|
|
@ -1,6 +1,7 @@
|
|||
"""简洁的RAG基础库"""
|
||||
|
||||
from .core import BaseRAG
|
||||
from .image_processor import ImageProcessor
|
||||
|
||||
__version__ = "0.1.0"
|
||||
__all__ = ["BaseRAG"]
|
||||
__all__ = ["BaseRAG", "ImageProcessor"]
|
||||
|
|
|
@ -337,6 +337,26 @@ class ModelManager:
|
|||
else:
|
||||
raise ValueError(f"不支持的重排模型类型: {config_type},支持的类型: 'local', 'api'")
|
||||
|
||||
@staticmethod
|
||||
def create_image_model(config: Dict) -> Any:
|
||||
"""创建图片处理模型(在线程池中运行)"""
|
||||
try:
|
||||
from .image_processor import ImageProcessor
|
||||
|
||||
config_type = config.get("type", "local")
|
||||
print(f"🖼️ 正在创建图片处理模型 ({config_type} 模式)...")
|
||||
|
||||
processor = ImageProcessor(config)
|
||||
print("✅ 图片处理模型创建成功")
|
||||
return processor
|
||||
|
||||
except ImportError:
|
||||
print("❌ 需要安装图片处理依赖: pip install transformers torch torchvision Pillow")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"❌ 图片处理模型创建失败: {e}")
|
||||
return None
|
||||
|
||||
|
||||
class BaseRAG(ABC):
|
||||
|
||||
|
@ -348,6 +368,7 @@ class BaseRAG(ABC):
|
|||
llm: Optional[BaseLLM] = None,
|
||||
embedding_config: Optional[Dict] = None,
|
||||
rerank_config: Optional[Dict] = None,
|
||||
image_config: Optional[Dict] = None,
|
||||
storage_directory: str = "./documents",
|
||||
status_db_path: str = "./file_status.db",
|
||||
):
|
||||
|
@ -359,6 +380,7 @@ class BaseRAG(ABC):
|
|||
:param llm: 可选的对话模型
|
||||
:param persist_directory: Chroma持久化目录
|
||||
:param rerank_config: 重排配置
|
||||
:param image_config: 图片处理配置
|
||||
:param storage_directory: 文件存储目录
|
||||
:param status_db_path: 文件状态数据库路径
|
||||
|
||||
|
@ -371,6 +393,12 @@ class BaseRAG(ABC):
|
|||
{"enabled": True, "type": "local", "model": "BAAI/bge-reranker-base", "top_k": 3}
|
||||
{"enabled": True, "type": "local", "model_path": "/path/to/your/rerank/model", "top_k": 3}
|
||||
{"enabled": True, "type": "api", "api_url": "http://localhost:8000/rerank", "model": "reranker-model", "api_key": "your-key", "top_k": 3}
|
||||
|
||||
image_config 示例:
|
||||
禁用图片处理: {"enabled": False}
|
||||
本地BLIP模型: {"enabled": True, "type": "local", "model": "Salesforce/blip-image-captioning-base"}
|
||||
本地模型路径: {"enabled": True, "type": "local", "model_path": "/path/to/your/image/model"}
|
||||
API图片处理: {"enabled": True, "type": "api", "api_url": "http://localhost:8000/image2text", "api_key": "your-key", "model": "image-caption"}
|
||||
"""
|
||||
self.vector_store_name = vector_store_name
|
||||
self.embedding_config = embedding_config or {
|
||||
|
@ -381,6 +409,7 @@ class BaseRAG(ABC):
|
|||
self.llm = llm
|
||||
self.persist_directory = persist_directory
|
||||
self.rerank_config = rerank_config or {"enabled": False}
|
||||
self.image_config = image_config or {"enabled": True}
|
||||
|
||||
# 初始化文件管理器
|
||||
self.file_manager = FileManager(storage_directory, status_db_path)
|
||||
|
@ -410,6 +439,13 @@ class BaseRAG(ABC):
|
|||
self.rerank_config, "rerank", ModelManager.create_rerank_model
|
||||
)
|
||||
|
||||
# 初始化图片处理模型
|
||||
self.image_processor = None
|
||||
if self.image_config.get("enabled", True):
|
||||
self.image_processor = await ModelManager.get_or_create_model(
|
||||
self.image_config, "image", ModelManager.create_image_model
|
||||
)
|
||||
|
||||
# 初始化 Chroma 向量库
|
||||
self.vector_store = Chroma(
|
||||
collection_name=self.vector_store_name,
|
||||
|
@ -543,6 +579,8 @@ class BaseRAG(ABC):
|
|||
"""
|
||||
根据文件类型异步加载文档
|
||||
"""
|
||||
await self._ensure_initialized() # 确保模型已初始化
|
||||
|
||||
file_path = Path(file_path)
|
||||
file_extension = file_path.suffix.lower()
|
||||
|
||||
|
@ -557,11 +595,39 @@ class BaseRAG(ABC):
|
|||
return loader.load()
|
||||
|
||||
elif file_extension in ['.doc', '.docx']:
|
||||
# Word文档
|
||||
# Word文档 - 增强图片处理
|
||||
try:
|
||||
from langchain_community.document_loaders import UnstructuredWordDocumentLoader
|
||||
from langchain_core.documents import Document
|
||||
|
||||
# 加载基本文档内容
|
||||
loader = UnstructuredWordDocumentLoader(str(file_path))
|
||||
return loader.load()
|
||||
documents = loader.load()
|
||||
|
||||
# 如果启用了图片处理,尝试提取图片
|
||||
if self.image_processor:
|
||||
try:
|
||||
from .image_processor import extract_images_from_docx
|
||||
images_info = extract_images_from_docx(str(file_path), self.image_processor)
|
||||
|
||||
if images_info:
|
||||
print(f"📸 从DOCX中提取到 {len(images_info)} 张图片")
|
||||
# 为每张图片创建单独的文档
|
||||
for image_path, description in images_info:
|
||||
image_doc = Document(
|
||||
page_content=description,
|
||||
metadata={
|
||||
"source": str(file_path),
|
||||
"type": "image",
|
||||
"image_path": image_path
|
||||
}
|
||||
)
|
||||
documents.append(image_doc)
|
||||
except Exception as e:
|
||||
print(f"图片提取失败,继续处理文本内容: {e}")
|
||||
|
||||
return documents
|
||||
|
||||
except ImportError:
|
||||
print("警告: 需要安装 unstructured 和 python-docx 来处理Word文档")
|
||||
print("请运行: pip install unstructured python-docx")
|
||||
|
@ -620,11 +686,39 @@ class BaseRAG(ABC):
|
|||
raise
|
||||
|
||||
elif file_extension == '.pdf':
|
||||
# PDF文件
|
||||
# PDF文件 - 增强图片处理
|
||||
try:
|
||||
from langchain_community.document_loaders import PyPDFLoader
|
||||
from langchain_core.documents import Document
|
||||
|
||||
# 加载基本PDF内容
|
||||
loader = PyPDFLoader(str(file_path))
|
||||
return loader.load()
|
||||
documents = loader.load()
|
||||
|
||||
# 如果启用了图片处理,尝试提取图片
|
||||
if self.image_processor:
|
||||
try:
|
||||
from .image_processor import extract_images_from_pdf
|
||||
images_info = extract_images_from_pdf(str(file_path), self.image_processor)
|
||||
|
||||
if images_info:
|
||||
print(f"📸 从PDF中提取到 {len(images_info)} 张图片")
|
||||
# 为每张图片创建单独的文档
|
||||
for image_path, description in images_info:
|
||||
image_doc = Document(
|
||||
page_content=description,
|
||||
metadata={
|
||||
"source": str(file_path),
|
||||
"type": "image",
|
||||
"image_path": image_path
|
||||
}
|
||||
)
|
||||
documents.append(image_doc)
|
||||
except Exception as e:
|
||||
print(f"PDF图片提取失败,继续处理文本内容: {e}")
|
||||
|
||||
return documents
|
||||
|
||||
except ImportError:
|
||||
try:
|
||||
# 备用方案:使用pdfplumber
|
||||
|
@ -640,6 +734,28 @@ class BaseRAG(ABC):
|
|||
page_content=text,
|
||||
metadata={"source": str(file_path), "page": i + 1}
|
||||
))
|
||||
|
||||
# 如果启用了图片处理,尝试提取图片
|
||||
if self.image_processor:
|
||||
try:
|
||||
from .image_processor import extract_images_from_pdf
|
||||
images_info = extract_images_from_pdf(str(file_path), self.image_processor)
|
||||
|
||||
if images_info:
|
||||
print(f"📸 从PDF中提取到 {len(images_info)} 张图片")
|
||||
for image_path, description in images_info:
|
||||
image_doc = Document(
|
||||
page_content=description,
|
||||
metadata={
|
||||
"source": str(file_path),
|
||||
"type": "image",
|
||||
"image_path": image_path
|
||||
}
|
||||
)
|
||||
documents.append(image_doc)
|
||||
except Exception as e:
|
||||
print(f"PDF图片提取失败: {e}")
|
||||
|
||||
return documents
|
||||
except ImportError:
|
||||
print("警告: 需要安装 PyPDF2 或 pdfplumber 来处理PDF文件")
|
||||
|
|
|
@ -0,0 +1,378 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
图片处理模块 - 简洁的图像到文本转换
|
||||
"""
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from typing import List, Dict, Optional, Tuple
|
||||
from PIL import Image
|
||||
|
||||
# 过滤警告
|
||||
warnings.filterwarnings("ignore", category=FutureWarning)
|
||||
warnings.filterwarnings("ignore", category=UserWarning)
|
||||
|
||||
|
||||
class ImageProcessor:
|
||||
"""图片处理器 - 支持多种配置方式的图像描述"""
|
||||
|
||||
def __init__(self, config: Dict = None):
|
||||
"""
|
||||
初始化图片处理器
|
||||
|
||||
Args:
|
||||
config: 配置字典,支持本地模型和API模式
|
||||
本地模型: {"type": "local", "model": "Salesforce/blip-image-captioning-base"}
|
||||
本地路径: {"type": "local", "model_path": "/path/to/model"}
|
||||
API调用: {"type": "api", "api_url": "http://localhost:8000/image2text", "api_key": "your-key"}
|
||||
"""
|
||||
self.config = config or {"type": "local", "model": "Salesforce/blip-image-captioning-base"}
|
||||
self.config_type = self.config.get("type", "local")
|
||||
self.model = None
|
||||
self.processor = None
|
||||
|
||||
def _load_model(self):
|
||||
"""根据配置加载模型"""
|
||||
if self.model is not None:
|
||||
return
|
||||
|
||||
if self.config_type == "local":
|
||||
self._load_local_model()
|
||||
elif self.config_type == "api":
|
||||
self._init_api_config()
|
||||
elif self.config_type == "basic":
|
||||
self._init_basic_config()
|
||||
else:
|
||||
raise ValueError(f"不支持的图片处理类型: {self.config_type},支持的类型: 'local', 'api', 'basic'")
|
||||
|
||||
def _load_local_model(self):
|
||||
"""加载本地模型"""
|
||||
try:
|
||||
from transformers import BlipProcessor, BlipForConditionalGeneration
|
||||
|
||||
# 支持本地路径和模型名称两种方式
|
||||
if "model_path" in self.config:
|
||||
model_name = self.config["model_path"]
|
||||
print(f"🖼️ 从本地路径加载图像模型: {model_name}")
|
||||
else:
|
||||
model_name = self.config.get("model", "Salesforce/blip-image-captioning-base")
|
||||
print(f"🖼️ 从HuggingFace Hub加载图像模型: {model_name}")
|
||||
|
||||
self.processor = BlipProcessor.from_pretrained(model_name)
|
||||
self.model = BlipForConditionalGeneration.from_pretrained(model_name)
|
||||
print("✅ 本地图像模型加载成功")
|
||||
|
||||
except ImportError:
|
||||
print("❌ 需要安装: pip install transformers torch torchvision")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"❌ 本地图像模型加载失败: {e}")
|
||||
raise
|
||||
|
||||
def _init_api_config(self):
|
||||
"""初始化API配置"""
|
||||
api_url = self.config.get("api_url")
|
||||
if not api_url:
|
||||
raise ValueError("使用API类型时必须提供api_url")
|
||||
|
||||
print(f"🖼️ 连接到图像处理API: {api_url}")
|
||||
self.api_config = {
|
||||
"api_url": api_url,
|
||||
"model": self.config.get("model", "image2text"),
|
||||
"api_key": self.config.get("api_key", "dummy"),
|
||||
"max_retries": self.config.get("max_retries", 3),
|
||||
}
|
||||
print("✅ API图像处理配置完成")
|
||||
|
||||
def _init_basic_config(self):
|
||||
"""初始化基础模式配置"""
|
||||
print("🖼️ 使用基础图片信息提取模式")
|
||||
self.basic_mode = True
|
||||
print("✅ 基础模式配置完成")
|
||||
|
||||
def extract_image_description(self, image_path: str) -> str:
|
||||
"""从图片提取文本描述"""
|
||||
try:
|
||||
self._load_model()
|
||||
|
||||
# 加载图片
|
||||
image = Image.open(image_path).convert('RGB')
|
||||
|
||||
if self.config_type == "local":
|
||||
return self._process_with_local_model(image)
|
||||
elif self.config_type == "api":
|
||||
return self._process_with_api(image_path, image)
|
||||
elif self.config_type == "basic":
|
||||
return self._basic_image_info(image_path, image)
|
||||
else:
|
||||
return self._basic_image_info(image_path, image)
|
||||
|
||||
except Exception as e:
|
||||
print(f"图片处理失败 {image_path}: {e}")
|
||||
return f"图片文件: {os.path.basename(image_path)} (处理失败)"
|
||||
|
||||
def _process_with_local_model(self, image: Image.Image) -> str:
|
||||
"""使用本地模型处理图片"""
|
||||
try:
|
||||
if self.model is None:
|
||||
return f"本地模型未加载"
|
||||
|
||||
inputs = self.processor(image, return_tensors="pt")
|
||||
out = self.model.generate(**inputs, max_length=50, num_beams=3)
|
||||
caption = self.processor.decode(out[0], skip_special_tokens=True)
|
||||
|
||||
return f"图片描述: {caption}"
|
||||
|
||||
except Exception as e:
|
||||
print(f"本地模型处理失败: {e}")
|
||||
return f"图片内容 (本地模型处理失败)"
|
||||
|
||||
def _process_with_api(self, image_path: str, image: Image.Image) -> str:
|
||||
"""使用API处理图片"""
|
||||
try:
|
||||
import base64
|
||||
import io
|
||||
import requests
|
||||
|
||||
# 将图片转换为base64
|
||||
buffered = io.BytesIO()
|
||||
image.save(buffered, format="JPEG")
|
||||
img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
|
||||
|
||||
# 准备API请求
|
||||
payload = {
|
||||
"model": self.api_config["model"],
|
||||
"image": img_base64,
|
||||
"format": "base64"
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self.api_config['api_key']}"
|
||||
}
|
||||
|
||||
# 发送请求
|
||||
response = requests.post(
|
||||
self.api_config["api_url"],
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=30
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
caption = result.get("description", result.get("caption", "API返回格式异常"))
|
||||
return f"图片描述: {caption}"
|
||||
else:
|
||||
return f"API调用失败: {response.status_code}"
|
||||
|
||||
except Exception as e:
|
||||
print(f"API处理失败: {e}")
|
||||
return f"图片内容 (API处理失败)"
|
||||
|
||||
def _basic_image_info(self, image_path: str, image: Image.Image) -> str:
|
||||
"""基础图片信息提取 - 增强版本,包含OCR文本提取"""
|
||||
filename = os.path.basename(image_path)
|
||||
width, height = image.size
|
||||
|
||||
# 尝试OCR文本提取
|
||||
ocr_text = self._extract_text_from_image(image)
|
||||
|
||||
# 基于文件名推测内容类型
|
||||
name_lower = filename.lower()
|
||||
if any(word in name_lower for word in ['python', 'py']):
|
||||
content_type = "Python编程相关图片"
|
||||
elif any(word in name_lower for word in ['chart', 'graph', 'data']):
|
||||
content_type = "图表或数据可视化"
|
||||
elif any(word in name_lower for word in ['diagram', 'flow', 'architecture']):
|
||||
content_type = "流程图或架构图"
|
||||
elif any(word in name_lower for word in ['ui', 'interface', 'screen']):
|
||||
content_type = "用户界面截图"
|
||||
else:
|
||||
content_type = "技术文档图片"
|
||||
|
||||
# 构建完整的图片描述
|
||||
description = f"图片文件: {filename} | 尺寸: {width}x{height} | 类型: {content_type}"
|
||||
|
||||
# 如果提取到文本,添加到描述中
|
||||
if ocr_text:
|
||||
description += f"\n📝 图片中的文本内容: {ocr_text}"
|
||||
|
||||
return description
|
||||
|
||||
def _extract_text_from_image(self, image: Image.Image) -> str:
|
||||
"""从图片中提取文本内容 (OCR)"""
|
||||
try:
|
||||
# 尝试使用pytesseract进行OCR
|
||||
import pytesseract
|
||||
|
||||
# 提取文本
|
||||
text = pytesseract.image_to_string(image, lang='eng+chi_sim')
|
||||
|
||||
# 清理和格式化文本
|
||||
if text:
|
||||
# 移除多余的空白字符
|
||||
lines = [line.strip() for line in text.split('\n') if line.strip()]
|
||||
cleaned_text = ' '.join(lines)
|
||||
|
||||
# 限制文本长度
|
||||
if len(cleaned_text) > 200:
|
||||
cleaned_text = cleaned_text[:200] + "..."
|
||||
|
||||
return cleaned_text
|
||||
|
||||
except ImportError:
|
||||
# 如果没有安装pytesseract,尝试使用easyocr
|
||||
try:
|
||||
import easyocr
|
||||
|
||||
# 创建OCR读取器(支持中英文)
|
||||
if not hasattr(self, '_ocr_reader'):
|
||||
self._ocr_reader = easyocr.Reader(['en', 'ch_sim'])
|
||||
|
||||
# 转换PIL图像为numpy数组
|
||||
import numpy as np
|
||||
img_array = np.array(image)
|
||||
|
||||
# 执行OCR
|
||||
results = self._ocr_reader.readtext(img_array)
|
||||
|
||||
# 提取文本
|
||||
if results:
|
||||
texts = [result[1] for result in results if result[2] > 0.5] # 置信度>0.5
|
||||
combined_text = ' '.join(texts)
|
||||
|
||||
# 限制文本长度
|
||||
if len(combined_text) > 200:
|
||||
combined_text = combined_text[:200] + "..."
|
||||
|
||||
return combined_text
|
||||
|
||||
except ImportError:
|
||||
# 如果都没有安装OCR库,返回提示
|
||||
return "(需要安装pytesseract或easyocr进行文字识别)"
|
||||
except Exception as e:
|
||||
print(f"OCR文本提取失败: {e}")
|
||||
return "(文字识别失败)"
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def extract_images_from_docx(docx_path: str, image_processor: ImageProcessor = None) -> List[Tuple[str, str]]:
|
||||
"""从DOCX文件中提取图片并生成描述"""
|
||||
try:
|
||||
from docx import Document
|
||||
|
||||
doc = Document(docx_path)
|
||||
images_info = []
|
||||
|
||||
# 使用传入的处理器或创建默认处理器
|
||||
processor = image_processor or ImageProcessor()
|
||||
|
||||
for rel in doc.part.rels.values():
|
||||
if "image" in rel.target_ref:
|
||||
image_data = rel.target_part.blob
|
||||
image_filename = rel.target_ref.split('/')[-1]
|
||||
|
||||
# 临时保存图片
|
||||
temp_path = f"/tmp/{image_filename}"
|
||||
with open(temp_path, 'wb') as f:
|
||||
f.write(image_data)
|
||||
|
||||
# 生成描述
|
||||
description = processor.extract_image_description(temp_path)
|
||||
images_info.append((temp_path, description))
|
||||
|
||||
# 清理临时文件
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
|
||||
return images_info
|
||||
|
||||
except Exception as e:
|
||||
print(f"DOCX图片提取失败: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def extract_images_from_pdf(pdf_path: str, image_processor: ImageProcessor = None) -> List[Tuple[str, str]]:
|
||||
"""从PDF文件中提取图片并生成描述"""
|
||||
try:
|
||||
import fitz # PyMuPDF
|
||||
|
||||
doc = fitz.open(pdf_path)
|
||||
images_info = []
|
||||
|
||||
# 使用传入的处理器或创建默认处理器
|
||||
processor = image_processor or ImageProcessor()
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
image_list = page.get_images()
|
||||
|
||||
for img_index, img in enumerate(image_list):
|
||||
xref = img[0]
|
||||
pix = fitz.Pixmap(doc, xref)
|
||||
|
||||
if pix.n - pix.alpha < 4: # RGB或灰度图
|
||||
img_filename = f"pdf_page_{page_num+1}_img_{img_index+1}.png"
|
||||
temp_path = f"/tmp/{img_filename}"
|
||||
pix.save(temp_path)
|
||||
|
||||
# 生成描述
|
||||
description = processor.extract_image_description(temp_path)
|
||||
images_info.append((temp_path, f"PDF第{page_num+1}页: {description}"))
|
||||
|
||||
# 清理临时文件
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
|
||||
pix = None
|
||||
|
||||
doc.close()
|
||||
return images_info
|
||||
|
||||
except Exception as e:
|
||||
print(f"PDF图片提取失败: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def extract_images_from_pdf(pdf_path: str, image_processor: 'ImageProcessor' = None) -> List[Tuple[str, str]]:
|
||||
"""从PDF文件中提取图片并生成描述"""
|
||||
try:
|
||||
import fitz # PyMuPDF
|
||||
|
||||
doc = fitz.open(pdf_path)
|
||||
images_info = []
|
||||
|
||||
# 使用传入的处理器或创建新的
|
||||
processor = image_processor or ImageProcessor()
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
image_list = page.get_images()
|
||||
|
||||
for img_index, img in enumerate(image_list):
|
||||
xref = img[0]
|
||||
pix = fitz.Pixmap(doc, xref)
|
||||
|
||||
if pix.n - pix.alpha < 4: # RGB或灰度图
|
||||
img_filename = f"pdf_page_{page_num+1}_img_{img_index+1}.png"
|
||||
temp_path = f"/tmp/{img_filename}"
|
||||
pix.save(temp_path)
|
||||
|
||||
# 生成描述
|
||||
description = processor.extract_image_description(temp_path)
|
||||
images_info.append((temp_path, f"PDF第{page_num+1}页: {description}"))
|
||||
|
||||
# 清理临时文件
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
|
||||
pix = None
|
||||
|
||||
doc.close()
|
||||
return images_info
|
||||
|
||||
except Exception as e:
|
||||
print(f"PDF图片提取失败: {e}")
|
||||
return []
|
Loading…
Reference in New Issue