feat: 支持ocr

This commit is contained in:
李如威 2025-08-08 17:10:11 +08:00
parent 5d50b9a23e
commit 1e2284728f
3 changed files with 500 additions and 5 deletions

View File

@ -1,6 +1,7 @@
"""简洁的RAG基础库""" """简洁的RAG基础库"""
from .core import BaseRAG from .core import BaseRAG
from .image_processor import ImageProcessor
__version__ = "0.1.0" __version__ = "0.1.0"
__all__ = ["BaseRAG"] __all__ = ["BaseRAG", "ImageProcessor"]

View File

@ -337,6 +337,26 @@ class ModelManager:
else: else:
raise ValueError(f"不支持的重排模型类型: {config_type},支持的类型: 'local', 'api'") 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): class BaseRAG(ABC):
@ -348,6 +368,7 @@ class BaseRAG(ABC):
llm: Optional[BaseLLM] = None, llm: Optional[BaseLLM] = None,
embedding_config: Optional[Dict] = None, embedding_config: Optional[Dict] = None,
rerank_config: Optional[Dict] = None, rerank_config: Optional[Dict] = None,
image_config: Optional[Dict] = None,
storage_directory: str = "./documents", storage_directory: str = "./documents",
status_db_path: str = "./file_status.db", status_db_path: str = "./file_status.db",
): ):
@ -359,6 +380,7 @@ class BaseRAG(ABC):
:param llm: 可选的对话模型 :param llm: 可选的对话模型
:param persist_directory: Chroma持久化目录 :param persist_directory: Chroma持久化目录
:param rerank_config: 重排配置 :param rerank_config: 重排配置
:param image_config: 图片处理配置
:param storage_directory: 文件存储目录 :param storage_directory: 文件存储目录
:param status_db_path: 文件状态数据库路径 :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": "BAAI/bge-reranker-base", "top_k": 3}
{"enabled": True, "type": "local", "model_path": "/path/to/your/rerank/model", "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} {"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.vector_store_name = vector_store_name
self.embedding_config = embedding_config or { self.embedding_config = embedding_config or {
@ -381,6 +409,7 @@ class BaseRAG(ABC):
self.llm = llm self.llm = llm
self.persist_directory = persist_directory self.persist_directory = persist_directory
self.rerank_config = rerank_config or {"enabled": False} 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) 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.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 向量库 # 初始化 Chroma 向量库
self.vector_store = Chroma( self.vector_store = Chroma(
collection_name=self.vector_store_name, collection_name=self.vector_store_name,
@ -543,6 +579,8 @@ class BaseRAG(ABC):
""" """
根据文件类型异步加载文档 根据文件类型异步加载文档
""" """
await self._ensure_initialized() # 确保模型已初始化
file_path = Path(file_path) file_path = Path(file_path)
file_extension = file_path.suffix.lower() file_extension = file_path.suffix.lower()
@ -557,11 +595,39 @@ class BaseRAG(ABC):
return loader.load() return loader.load()
elif file_extension in ['.doc', '.docx']: elif file_extension in ['.doc', '.docx']:
# Word文档 # Word文档 - 增强图片处理
try: try:
from langchain_community.document_loaders import UnstructuredWordDocumentLoader from langchain_community.document_loaders import UnstructuredWordDocumentLoader
from langchain_core.documents import Document
# 加载基本文档内容
loader = UnstructuredWordDocumentLoader(str(file_path)) 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: except ImportError:
print("警告: 需要安装 unstructured 和 python-docx 来处理Word文档") print("警告: 需要安装 unstructured 和 python-docx 来处理Word文档")
print("请运行: pip install unstructured python-docx") print("请运行: pip install unstructured python-docx")
@ -620,11 +686,39 @@ class BaseRAG(ABC):
raise raise
elif file_extension == '.pdf': elif file_extension == '.pdf':
# PDF文件 # PDF文件 - 增强图片处理
try: try:
from langchain_community.document_loaders import PyPDFLoader from langchain_community.document_loaders import PyPDFLoader
from langchain_core.documents import Document
# 加载基本PDF内容
loader = PyPDFLoader(str(file_path)) 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: except ImportError:
try: try:
# 备用方案使用pdfplumber # 备用方案使用pdfplumber
@ -640,6 +734,28 @@ class BaseRAG(ABC):
page_content=text, page_content=text,
metadata={"source": str(file_path), "page": i + 1} 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 return documents
except ImportError: except ImportError:
print("警告: 需要安装 PyPDF2 或 pdfplumber 来处理PDF文件") print("警告: 需要安装 PyPDF2 或 pdfplumber 来处理PDF文件")

View File

@ -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 []