From 1e2284728f3aa5d468acd82b9edfddedb8f92e45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Fri, 8 Aug 2025 17:10:11 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81ocr?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/base_rag/__init__.py | 3 +- src/base_rag/core.py | 124 ++++++++++- src/base_rag/image_processor.py | 378 ++++++++++++++++++++++++++++++++ 3 files changed, 500 insertions(+), 5 deletions(-) create mode 100644 src/base_rag/image_processor.py diff --git a/src/base_rag/__init__.py b/src/base_rag/__init__.py index b1f1647..cbab32f 100644 --- a/src/base_rag/__init__.py +++ b/src/base_rag/__init__.py @@ -1,6 +1,7 @@ """简洁的RAG基础库""" from .core import BaseRAG +from .image_processor import ImageProcessor __version__ = "0.1.0" -__all__ = ["BaseRAG"] +__all__ = ["BaseRAG", "ImageProcessor"] diff --git a/src/base_rag/core.py b/src/base_rag/core.py index e1c1697..8530603 100644 --- a/src/base_rag/core.py +++ b/src/base_rag/core.py @@ -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文件") diff --git a/src/base_rag/image_processor.py b/src/base_rag/image_processor.py new file mode 100644 index 0000000..26fa0f1 --- /dev/null +++ b/src/base_rag/image_processor.py @@ -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 []