import os import uuid from pathlib import Path from argparse import ArgumentParser import fitz # fitz就是pip install PyMuPDF import numpy as np # Added import for numpy from PIL import Image from tqdm import tqdm from rapidocr_onnxruntime import RapidOCR from src.utils import logger, is_text_pdf GOLBAL_STATE = {} class OCRPlugin: """OCR 插件""" def __init__(self, **kwargs): self.ocr = None self.det_box_thresh = kwargs.get('det_box_thresh', 0.3) def load_model(self): """加载 OCR 模型""" logger.info(f"加载 OCR 模型，仅在第一次调用时加载") model_dir = os.path.join(os.getenv("MODEL_DIR", ""), "SWHL/RapidOCR") det_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_det_infer.onnx") rec_model_dir = os.path.join(model_dir, "PP-OCRv4/ch_PP-OCRv4_rec_infer.onnx") assert os.path.exists(model_dir), ( f"模型文件不存在，请下载 SWHL/RapidOCR 到 {model_dir}，" "并确认是否在 docker-compose.dev.yml 中添加 MODEL_DIR 环境变量" ) self.ocr = RapidOCR(det_box_thresh=0.3, det_model_path=det_model_dir, rec_model_path=rec_model_dir) logger.info(f"OCR Plugin for det_box_thresh = {self.det_box_thresh} loaded.") def process_image(self, image): """ 对单张图像执行OCR并提取文本 Args: image: 图像数据，支持多种格式： - str: 图像文件路径 - PIL.Image: PIL图像对象 - numpy.ndarray: numpy图像数组 Returns: str: 提取的文本内容 """ # 确保模型已加载 if self.ocr is None: self.load_model() # 处理不同类型的输入图像 try: if isinstance(image, str): # 图像路径直接传递给OCR处理 image_path = image is_temp_file = False else: # 创建临时文件 is_temp_file = True image_path = self._create_temp_image_file(image) # 执行 OCR result, _ = self.ocr(image_path) # 清理临时文件 if is_temp_file and os.path.exists(image_path): os.remove(image_path) # 提取文本 if result: text = '\n'.join([line[1] for line in result]) return text else: logger.warning(f"OCR未能识别出文本内容") return "" except Exception as e: logger.error(f"OCR处理失败: {str(e)}") raise def _create_temp_image_file(self, image): """ 将图像数据保存为临时文件 Args: image: PIL.Image或numpy.ndarray格式的图像数据 Returns: str: 临时文件路径 """ # 为临时文件创建目录（如果不存在） tmp_dir = os.path.join(os.getcwd(), 'tmp') os.makedirs(tmp_dir, exist_ok=True) # 生成临时文件路径 temp_filename = f'ocr_temp_{uuid.uuid4().hex[:8]}.png' image_path = os.path.join(tmp_dir, temp_filename) # 根据图像类型保存文件 if isinstance(image, Image.Image): # 保存PIL图像对象到临时文件 image.save(image_path) elif isinstance(image, np.ndarray): # 将numpy数组转换为PIL图像并保存 Image.fromarray(image).save(image_path) else: raise ValueError("不支持的图像类型，必须是PIL.Image或numpy数组") return image_path def process_pdf(self, pdf_path): """ 处理PDF文件并提取文本 :param pdf_path: PDF文件路径 :return: 提取的文本 """ if not os.path.exists(pdf_path): raise FileNotFoundError(f"PDF file not found: {pdf_path}") try: # 检查是否为文本PDF，可能会出现错误，比如每一页都有可读取的水印文字，但是内容本身是扫描件，需要使用OCR处理 # if is_text_pdf(pdf_path): # from src.core.indexing import pdfreader # logger.info("PDF file is text, use llama_index.readers.file to read") # return pdfreader(pdf_path) images = [] pdfDoc = fitz.open(pdf_path) totalPage = pdfDoc.page_count for pg in tqdm(range(totalPage), desc='to images', ncols=100): page = pdfDoc[pg] rotate, zoom_x, zoom_y = 0, 2, 2 mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate) pix = page.get_pixmap(matrix=mat, alpha=False) img_pil = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) images.append(img_pil) # 处理每个图像并合并文本 all_text = [] for img_path in tqdm(images, desc='to txt', ncols=100): text = self.process_image(img_path) all_text.append(text) return '\n\n'.join(all_text) except Exception as e: logger.error(f"PDF processing error: {str(e)}") return "" def process_pdf_mineru(self, pdf_path): """ 使用Mineru OCR处理PDF文件 :param pdf_path: PDF文件路径 :return: 提取的文本 """ mineru_ocr_uri = os.getenv("MINERU_OCR_URI", "http://localhost:5051") import requests import json health_check_response = requests.get(f"{mineru_ocr_uri}/health", timeout=5) if health_check_response.status_code != 200 or health_check_response.json().get("status") != "healthy": logger.error("Mineru OCR service health check failed.") raise RuntimeError("Mineru OCR service health check failed. Please check the log use `docker logs mineru-api`") # 读取PDF文件 with open(pdf_path, 'rb') as f: files = {'file': f} data = { 'parse_method': 'ocr', # 使用OCR模式 'is_json_md_dump': False, # 不需要保存中间文件 'return_layout': False, # 不需要返回布局信息 'return_info': False, # 不需要返回额外信息 'return_content_list': False, # 不需要返回内容列表 'return_images': False, # 不需要返回图片 } try: # 发送POST请求到Mineru OCR服务 response = requests.post( f"{mineru_ocr_uri}/file_parse", files=files, data=data ) response.raise_for_status() # 检查响应状态 # 解析响应 result = response.json() if 'md_content' in result: return result['md_content'] else: logger.error("Mineru OCR response does not contain md_content") return "" except requests.exceptions.RequestException as e: logger.error(f"Mineru OCR request failed: {str(e)}") return "" except json.JSONDecodeError as e: logger.error(f"Failed to parse Mineru OCR response: {str(e)}") return "" except Exception as e: logger.error(f"Unexpected error in Mineru OCR processing: {str(e)}") return "" def get_state(task_id): return GOLBAL_STATE.get(task_id, {}) def plainreader(file_path): """读取普通文本文件并返回text文本""" assert os.path.exists(file_path), "File not found" with open(file_path, "r") as f: text = f.read() return text if __name__ == "__main__": parser = ArgumentParser() parser.add_argument('--pdf-path', type=str, required=True, help='Path to the PDF file') parser.add_argument('--return-text', action='store_true', help='Return the extracted text') args = parser.parse_args() ocr = OCRPlugin() text = ocr.process_pdf(args.pdf_path) print(text)