feat: 支持ocr

This commit is contained in:
李如威 2025-08-08 18:08:11 +08:00
parent 1e2284728f
commit 3e4df252da
11 changed files with 312 additions and 254 deletions

100
demo_image_integration.py Normal file
View File

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

@ -30,12 +30,16 @@ class ImageProcessor:
self.config_type = self.config.get("type", "local")
self.model = None
self.processor = None
self.ocr_reader = None # EasyOCR读取器
def _load_model(self):
"""根据配置加载模型"""
"""根据配置加载模型和OCR"""
if self.model is not None:
return
# 初始化EasyOCR读取器
self._init_ocr_reader()
if self.config_type == "local":
self._load_local_model()
elif self.config_type == "api":
@ -45,6 +49,21 @@ class ImageProcessor:
else:
raise ValueError(f"不支持的图片处理类型: {self.config_type},支持的类型: 'local', 'api', 'basic'")
def _init_ocr_reader(self):
"""初始化EasyOCR读取器"""
try:
import easyocr
if self.ocr_reader is None:
print("🔍 正在初始化EasyOCR读取器...")
self.ocr_reader = easyocr.Reader(['en', 'ch_sim'])
print("✅ EasyOCR读取器初始化成功")
except ImportError:
print("⚠️ 未安装EasyOCROCR功能将受限: pip install easyocr")
self.ocr_reader = None
except Exception as e:
print(f"⚠️ EasyOCR初始化失败: {e}")
self.ocr_reader = None
def _load_local_model(self):
"""加载本地模型"""
try:
@ -90,6 +109,41 @@ class ImageProcessor:
self.basic_mode = True
print("✅ 基础模式配置完成")
def _extract_text_with_easyocr(self, image: Image.Image) -> str:
"""使用EasyOCR提取图片中的文本"""
try:
if self.ocr_reader is None:
return ""
# 转换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.6] # 置信度>0.6
combined_text = ' '.join(texts)
# 清理和格式化文本
if combined_text:
lines = [line.strip() for line in combined_text.split('\n') if line.strip()]
cleaned_text = ' '.join(lines)
# 限制文本长度
if len(cleaned_text) > 300:
cleaned_text = cleaned_text[:300] + "..."
return cleaned_text
except Exception as e:
print(f"EasyOCR文本提取失败: {e}")
return ""
return ""
def extract_image_description(self, image_path: str) -> str:
"""从图片提取文本描述"""
try:
@ -112,24 +166,42 @@ class ImageProcessor:
return f"图片文件: {os.path.basename(image_path)} (处理失败)"
def _process_with_local_model(self, image: Image.Image) -> str:
"""使用本地模型处理图片"""
"""使用本地模型处理图片结合OCR文本"""
try:
if self.model is None:
return f"本地模型未加载"
# 首先提取OCR文本
ocr_text = self._extract_text_with_easyocr(image)
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)
# 生成图片描述
description = ""
if self.model is not None:
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)
description = f"图片描述: {caption}"
else:
description = "本地模型未加载"
return f"图片描述: {caption}"
# 结合OCR文本和图片描述
if ocr_text:
return f"{description}\n📝 图片中的文字: {ocr_text}"
else:
return description
except Exception as e:
print(f"本地模型处理失败: {e}")
# 即使模型失败也尝试返回OCR文本
ocr_text = self._extract_text_with_easyocr(image)
if ocr_text:
return f"图片内容 (模型处理失败)\n📝 图片中的文字: {ocr_text}"
return f"图片内容 (本地模型处理失败)"
def _process_with_api(self, image_path: str, image: Image.Image) -> str:
"""使用API处理图片"""
"""使用API处理图片结合OCR文本"""
try:
# 首先提取OCR文本
ocr_text = self._extract_text_with_easyocr(image)
# 调用API生成图片描述
import base64
import io
import requests
@ -159,24 +231,36 @@ class ImageProcessor:
timeout=30
)
# 处理API响应
description = ""
if response.status_code == 200:
result = response.json()
caption = result.get("description", result.get("caption", "API返回格式异常"))
return f"图片描述: {caption}"
description = f"图片描述: {caption}"
else:
return f"API调用失败: {response.status_code}"
description = f"API调用失败: {response.status_code}"
# 结合OCR文本和API描述
if ocr_text:
return f"{description}\n📝 图片中的文字: {ocr_text}"
else:
return description
except Exception as e:
print(f"API处理失败: {e}")
# 即使API失败也尝试返回OCR文本
ocr_text = self._extract_text_with_easyocr(image)
if ocr_text:
return f"图片内容 (API处理失败)\n📝 图片中的文字: {ocr_text}"
return f"图片内容 (API处理失败)"
def _basic_image_info(self, image_path: str, image: Image.Image) -> str:
"""基础图片信息提取 - 增强版本包含OCR文本提取"""
"""基础图片信息提取 - 主要使用EasyOCR"""
filename = os.path.basename(image_path)
width, height = image.size
# 尝试OCR文本提取
ocr_text = self._extract_text_from_image(image)
# 使用EasyOCR提取文本
ocr_text = self._extract_text_with_easyocr(image)
# 基于文件名推测内容类型
name_lower = filename.lower()
@ -191,75 +275,20 @@ class ImageProcessor:
else:
content_type = "技术文档图片"
# 构建完整的图片描述
# 构建图片描述
description = f"图片文件: {filename} | 尺寸: {width}x{height} | 类型: {content_type}"
# 如果提取到文本,添加到描述中
if ocr_text:
description += f"\n📝 图片中的文本内容: {ocr_text}"
description += f"\n📝 图片中的文字: {ocr_text}"
else:
description += "\n📝 未检测到文字内容"
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文件中提取图片并生成描述"""
"""从DOCX文件中提取图片并进行文字识别和内容分析"""
try:
from docx import Document
@ -279,7 +308,7 @@ def extract_images_from_docx(docx_path: str, image_processor: ImageProcessor = N
with open(temp_path, 'wb') as f:
f.write(image_data)
# 生成描述
# 进行文字识别和内容分析
description = processor.extract_image_description(temp_path)
images_info.append((temp_path, description))
@ -295,7 +324,7 @@ def extract_images_from_docx(docx_path: str, image_processor: ImageProcessor = N
def extract_images_from_pdf(pdf_path: str, image_processor: ImageProcessor = None) -> List[Tuple[str, str]]:
"""从PDF文件中提取图片并生成描述"""
"""从PDF文件中提取图片并进行文字识别和内容分析"""
try:
import fitz # PyMuPDF
@ -318,49 +347,7 @@ def extract_images_from_pdf(pdf_path: str, image_processor: ImageProcessor = Non
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}"))

View File

@ -1,5 +0,0 @@
NumPy是Python中用于科学计算的基础库提供多维数组对象。
Pandas是强大的数据分析和处理库提供DataFrame数据结构。
Matplotlib是Python的绘图库用于创建静态、动态和交互式图表。
Scikit-learn是机器学习库提供各种算法和工具。

View File

@ -1,6 +0,0 @@
Python是一种高级编程语言。
它具有简洁的语法和强大的功能。
Python广泛应用于Web开发、数据科学、人工智能等领域。
机器学习库如scikit-learn、TensorFlow和PyTorch都支持Python。
Flask和Django是流行的Python Web框架。

View File

@ -1,44 +0,0 @@
# 机器学习入门
## 什么是机器学习?
机器学习是人工智能的一个分支,它使计算机能够在没有明确编程的情况下学习和改进。
## 主要类型
### 监督学习
- **分类**: 预测类别标签
- **回归**: 预测连续数值
### 无监督学习
- **聚类**: 发现数据中的群组
- **降维**: 减少特征数量
### 强化学习
- 通过与环境交互学习最优策略
## 常用算法
1. **线性回归**: 预测连续值
2. **逻辑回归**: 二分类问题
3. **决策树**: 易于理解和解释
4. **随机森林**: 集成学习方法
5. **支持向量机**: 处理高维数据
6. **神经网络**: 深度学习基础
## Python机器学习库
- **Scikit-learn**: 经典机器学习算法
- **TensorFlow**: 深度学习框架
- **PyTorch**: 动态深度学习框架
- **XGBoost**: 梯度提升算法
## 学习路径
1. 掌握Python基础
2. 学习数据处理(Pandas, NumPy)
3. 理解统计学基础
4. 实践经典算法
5. 深入深度学习
机器学习正在改变世界,值得每个人学习!

View File

@ -1,5 +0,0 @@
Python是一种高级编程语言由Guido van Rossum于1991年创建。
Python具有简洁易读的语法适合初学者学习编程。
Python是解释型语言支持面向对象、函数式等多种编程范式。
Python的设计哲学强调代码的可读性和简洁性。

View File

@ -1,70 +0,0 @@
# Python编程指南
## 基础语法
Python是一种高级编程语言以其简洁明了的语法而闻名。
### 变量和数据类型
```python
# 字符串
name = "Python"
# 整数
age = 30
# 浮点数
pi = 3.14159
# 布尔值
is_programming = True
```
### 控制结构
#### 条件语句
```python
if age >= 18:
print("成年人")
else:
print("未成年人")
```
#### 循环
```python
for i in range(5):
print(f"数字: {i}")
while count > 0:
print(count)
count -= 1
```
## 函数定义
```python
def greet(name):
return f"Hello, {name}!"
def calculate_area(radius):
return 3.14159 * radius ** 2
```
## 面向对象编程
```python
class Person:
def __init__(self, name, age):
self.name = name
self.age = age
def introduce(self):
return f"我是{self.name},今年{self.age}岁"
```
## 常用库
- **NumPy**: 科学计算
- **Pandas**: 数据分析
- **Matplotlib**: 数据可视化
- **Requests**: HTTP请求
- **Flask/Django**: Web开发
Python是学习编程的绝佳选择适合初学者入门。

View File

@ -1,5 +0,0 @@
Flask是一个轻量级的Python Web框架易于学习和使用。
Django是一个功能丰富的Python Web框架适合大型项目开发。
FastAPI是现代的Python Web框架专为构建API而设计。
Tornado是一个可扩展的非阻塞Web服务器和Web应用框架。

106
test_image_ocr.py Normal file
View File

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