import os import base64 import fitz # PyMuPDF import pdfplumber import win32com.client as win32 from docx import Document from openai import OpenAI from typing import List, Union import config.config class DocumentReader: def __init__(self, api_key: str = None, pdf_sample_pages: int = 2): """ :param api_key: Moonshot API Key :param pdf_sample_pages: 判断PDF是否为扫描件时抽样的页数（默认为2） """ self.client = OpenAI( api_key=api_key or config.config.MOONSHOT_API_KEY, base_url="https://api.moonshot.cn/v1", ) self.pdf_sample_pages = pdf_sample_pages # ================= PDF 处理（智能判断类型）================= def _read_pdf_with_validation(self, pdf_path: str) -> str: """ 读取PDF文件： - 抽样前 self.pdf_sample_pages 页，若均无文本则判定为扫描件，全部页面直接OCR； - 否则采用混合模式：每页优先提取文本，无文本时再OCR。 """ full_text = [] try: with fitz.open(pdf_path) as doc_fitz, pdfplumber.open(pdf_path) as doc_plumber: total_pages = len(doc_fitz) if total_pages == 0: return "[Error] PDF文件为空" # 抽样判断PDF类型 sample_limit = min(self.pdf_sample_pages, total_pages) has_text_page = False for page_num in range(sample_limit): # 仅用fitz快速判断是否有文本（pdfplumber较慢，这里只用fitz） text_fitz = doc_fitz[page_num].get_text().strip() if text_fitz: has_text_page = True break # 根据类型选择处理模式 if not has_text_page: # 纯扫描件：所有页直接OCR print(f"[Info] 检测为扫描型PDF，共{total_pages}页，直接进行OCR...") for page_num in range(total_pages): print(f"[Info] 正在OCR第 {page_num+1}/{total_pages} 页...") pix = doc_fitz[page_num].get_pixmap(dpi=150) img_bytes = pix.tobytes("png") ocr_result = self._read_images_via_kimi( [img_bytes], prompt="请提取图片中的全部文字内容，保留原始格式。" ) if ocr_result.startswith("[Error]"): page_content = f"[OCR失败] {ocr_result}" else: page_content = ocr_result full_text.append(f"--- Page {page_num + 1} ---\n{page_content}") else: # 混合型或文本型：逐页先提取文本，空则OCR print(f"[Info] 检测为文本型或混合型PDF，逐页处理...") for page_num in range(total_pages): # 尝试提取文本 text_fitz = doc_fitz[page_num].get_text().strip() text_plumber = doc_plumber.pages[page_num].extract_text() or "" page_content = text_plumber if len(text_plumber) > len(text_fitz) else text_fitz if not page_content.strip(): # 扫描页，进行OCR print(f"[Info] 第 {page_num+1} 页为扫描页，正在进行OCR...") pix = doc_fitz[page_num].get_pixmap(dpi=150) img_bytes = pix.tobytes("png") ocr_result = self._read_images_via_kimi( [img_bytes], prompt="请提取图片中的全部文字内容，保留原始格式。" ) if not ocr_result.startswith("[Error]"): page_content = ocr_result else: page_content = f"[OCR失败] {ocr_result}" full_text.append(f"--- Page {page_num + 1} ---\n{page_content}") except Exception as e: return f"[Error] PDF读取失败: {str(e)}" return "\n".join(full_text) # ================= Word 处理 ================= def _read_word(self, word_path: str) -> str: abs_path = os.path.abspath(word_path) ext = os.path.splitext(abs_path)[1].lower() if ext == ".docx": return self._extract_docx_structured_text(abs_path) elif ext == ".doc": temp_docx = abs_path + "x" word_app = win32.Dispatch('Word.Application') word_app.Visible = False try: doc = word_app.Documents.Open(abs_path) doc.SaveAs(temp_docx, FileFormat=16) # 16 = wdFormatXMLDocument doc.Close() text = self._extract_docx_structured_text(temp_docx) if os.path.exists(temp_docx): os.remove(temp_docx) return text except Exception as e: return f"[Error] .doc转换失败: {str(e)}" finally: word_app.Quit() return "[Error] 不支持的Word格式" def _extract_docx_structured_text(self, docx_path: str) -> str: """提取.docx中的段落和表格文本""" try: doc = Document(docx_path) full_content = [] for element in doc.element.body.iterchildren(): # 段落 if element.tag.endswith('p'): texts = [node.text for node in element.iter() if node.tag.endswith('t') and node.text] para_text = "".join(texts).strip() if para_text: full_content.append(para_text) # 表格 elif element.tag.endswith('tbl'): for table in doc.tables: if table._element == element: for row in table.rows: row_data = [cell.text.strip() for cell in row.cells if cell.text.strip()] if row_data: full_content.append(" | ".join(row_data)) break return "\n".join(full_content) except Exception as e: return f"[Error] Docx解析异常: {str(e)}" # ================= 图片 OCR（支持文件路径、base64 data URL、字节数据）================= def _read_images_via_kimi(self, image_inputs: List[Union[str, bytes]], prompt: str = "请提取图片中的文字内容") -> str: """ 通用图片OCR接口： - image_inputs 列表中的每个元素可以是： * 图片文件路径 (str) * data URL 字符串 (以 "data:image/" 开头) * 图片原始字节数据 (bytes) """ content_list = [] for img_input in image_inputs: if isinstance(img_input, str) and os.path.exists(img_input): # 文件路径 with open(img_input, "rb") as f: image_data = f.read() ext = os.path.splitext(img_input)[1].replace(".", "").lower() if ext not in ['png', 'jpg', 'jpeg', 'gif', 'bmp']: ext = 'png' # 默认 image_url = f"data:image/{ext};base64,{base64.b64encode(image_data).decode('utf-8')}" elif isinstance(img_input, str) and img_input.startswith("data:image/"): # 已经是 data URL image_url = img_input elif isinstance(img_input, bytes): # 字节数据，默认使用PNG格式 image_data = img_input image_url = f"data:image/png;base64,{base64.b64encode(image_data).decode('utf-8')}" else: continue # 忽略无效输入 content_list.append({"type": "image_url", "image_url": {"url": image_url}}) if not content_list: return "[Error] 没有有效的图片输入" content_list.append({"type": "text", "text": prompt}) try: completion = self.client.chat.completions.create( model="moonshot-v1-8k-vision-preview", messages=[{"role": "user", "content": content_list}], ) return completion.choices[0].message.content except Exception as e: return f"[Error] Kimi API失败: {str(e)}" # ================= 主接口 ================= def process_input(self, inputs: Union[str, List[str]]) -> str: """ 统一入口： - 如果 inputs 是列表，则视为多张图片路径，调用图片OCR。 - 如果是单个字符串，判断路径类型： * 目录不存在 → 报错 * PDF → 调用PDF处理（自动判断文本/扫描） * Word文档 → 调用Word处理 * 图片文件 → 调用图片OCR * 其他 → 报错 """ if isinstance(inputs, list): return self._read_images_via_kimi(inputs) if not os.path.exists(inputs): return f"[Error] 路径不存在: {inputs}" ext = os.path.splitext(inputs)[1].lower() if ext in [".pdf",".PDF"]: return self._read_pdf_with_validation(inputs) elif ext in [".docx", ".doc"]: return self._read_word(inputs) elif ext in [".png", ".jpg", ".jpeg", ".gif", ".bmp"]: return self._read_images_via_kimi([inputs]) return f"[Error] 未知文件格式: {ext}" # ================= 测试运行 ================= if __name__ == "__main__": # 请替换为你的 Moonshot API Key reader = DocumentReader() # 示例：处理一个可能为扫描件的PDF # test_path1 = r"E:\project\Ruling_and_Judgment_Comparative_Analysis_System\demo\（2024-3853）朱焱诉苏州星际浩宸企业管理有限公司（工资差额、经济补偿）.doc" # test_path2 = r"E:\project\Ruling_and_Judgment_Comparative_Analysis_System\demo\苏园劳人仲案字〔2024〕第3853号.pdf" # test_path3 = r"E:\project\Ruling_and_Judgment_Comparative_Analysis_System\demo\微信图片_20250925131023_134_4.jpg" test_path4 = r"E:\project\arbitration_system\transcript_extractor\test\庭审笔录.PDF" print(f"[Log] 开始处理: {test_path4}") result = reader.process_input(test_path4) print("\n--- 提取结果 ---\n") print(result)