feat: 支持ocr
This commit is contained in:
parent
5d50b9a23e
commit
1e2284728f
|
@ -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