diff --git a/requirements.txt b/requirements.txt index ea487e2..fc5bfab 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,7 @@ pytest-asyncio pymupdf python-docx scikit-learn -aiofiles \ No newline at end of file +aiofiles +paddlepaddle +paddleocr +pillow \ No newline at end of file diff --git a/src/pipeline/core/utils.py b/src/pipeline/core/utils.py index 0930e96..112af76 100644 --- a/src/pipeline/core/utils.py +++ b/src/pipeline/core/utils.py @@ -1,25 +1,55 @@ import asyncio import docx import fitz # PyMuPDF -from pathlib import Path import aiofiles +import io +from pathlib import Path +# from paddleocr import PaddleOCR +from PIL import Image # ----------------------------- -# 文件类型加载器 +# 文本加载 # ----------------------------- +# 初始化 OCR(只初始化一次),注意这是 CPU 版本,如果需要 GPU 需要额外配置 +print("ocr 加载...") +ocr = {} +# PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False) +print("ocr 加载结束") -async def load_txt(path: str) -> str: +async def _ocr_image_bytes(img_bytes: bytes) -> str: + """对图片字节流做 OCR(线程池避免阻塞 asyncio)""" + + def _ocr(): + img = Image.open(io.BytesIO(img_bytes)) + result = ocr.ocr(img, cls=True) + + # result: [[ ['text', 'score'], ... ]] + if not result or not result[0]: + return "" + + lines = [] + for line in result: + txt = line[1][0] # OCR 文字 + confidence = line[1][1] + lines.append(txt) + + return "\n".join(lines) + + return await asyncio.to_thread(_ocr) + + +async def _load_txt(path: str) -> str: async with aiofiles.open(path, "r", encoding="utf-8") as f: return await f.read() -async def load_md(path: str) -> str: +async def _load_md(path: str) -> str: async with aiofiles.open(path, "r", encoding="utf-8") as f: return await f.read() -async def load_docx(path: str) -> str: +async def _load_docx(path: str) -> str: # docx 读取是阻塞 I/O → 放入线程池 def _read(): doc = docx.Document(path) @@ -28,29 +58,56 @@ async def load_docx(path: str) -> str: return await asyncio.to_thread(_read) -async def load_pdf(path: str) -> str: - # PyMuPDF 也是阻塞 → 放入线程池 - def _read(): - pdf = fitz.open(path) - return "\n".join(page.get_text() for page in pdf) +async def _load_pdf(path: str) -> str: - return await asyncio.to_thread(_read) + def _read_pdf(): + doc = fitz.open(path) + pages = [] + for page in doc: + page_text = page.get_text() + # 获取图片 + image_bytes_list = [] + for img in page.get_images(full=True): + xref = img[0] + pix = fitz.Pixmap(doc, xref) + img_bytes = pix.tobytes("png") + image_bytes_list.append(img_bytes) -# ----------------------------- -# 统一调度器 -# ----------------------------- + pages.append((page_text, image_bytes_list)) + return pages + + # PDF 解析在子线程执行 + pages = await asyncio.to_thread(_read_pdf) + + final_text = [] + + # 分页处理 + OCR 并发 + for text, image_bytes_list in pages: + final_text.append(text) + + # 并发 OCR + ocr_tasks = [asyncio.create_task(_ocr_image_bytes(b)) for b in image_bytes_list] + + if ocr_tasks: + ocr_results = await asyncio.gather(*ocr_tasks) + final_text.extend(ocr_results) + + return "\n".join(final_text) async def load_document(path: str) -> str: suffix = Path(path).suffix.lower() + + print(f"读取文件: {path}") + if suffix == ".txt": - return await load_txt(path) + return await _load_txt(path) if suffix in (".md", ".markdown"): - return await load_md(path) + return await _load_md(path) if suffix == ".docx": - return await load_docx(path) + return await _load_docx(path) if suffix == ".pdf": - return await load_pdf(path) + return await _load_pdf(path) raise ValueError(f"Unsupported file type: {suffix}") diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py index 4e97cbb..91e0a5a 100644 --- a/src/tests/test_nodes.py +++ b/src/tests/test_nodes.py @@ -10,6 +10,7 @@ async def test_embedding(): shared = { "files": [ + "./files/带图片的.pdf", "./files/大白智问-API接入文档-V1.2.2.pdf", "./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md", "./files/deepsearch状态.txt",