feat: ocr
This commit is contained in:
parent
bc29326c2c
commit
bbd7c23fff
|
|
@ -7,4 +7,7 @@ pytest-asyncio
|
|||
pymupdf
|
||||
python-docx
|
||||
scikit-learn
|
||||
aiofiles
|
||||
aiofiles
|
||||
paddlepaddle
|
||||
paddleocr
|
||||
pillow
|
||||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ async def test_embedding():
|
|||
|
||||
shared = {
|
||||
"files": [
|
||||
"./files/带图片的.pdf",
|
||||
"./files/大白智问-API接入文档-V1.2.2.pdf",
|
||||
"./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md",
|
||||
"./files/deepsearch状态.txt",
|
||||
|
|
|
|||
Loading…
Reference in New Issue