feat: ocr

This commit is contained in:
李如威 2025-12-09 17:49:35 +08:00
parent bc29326c2c
commit bbd7c23fff
3 changed files with 80 additions and 19 deletions

View File

@ -7,4 +7,7 @@ pytest-asyncio
pymupdf
python-docx
scikit-learn
aiofiles
aiofiles
paddlepaddle
paddleocr
pillow

View File

@ -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}")

View File

@ -10,6 +10,7 @@ async def test_embedding():
shared = {
"files": [
"./files/带图片的.pdf",
"./files/大白智问-API接入文档-V1.2.2.pdf",
"./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md",
"./files/deepsearch状态.txt",