feat: ocr
This commit is contained in:
parent
bc29326c2c
commit
bbd7c23fff
|
|
@ -7,4 +7,7 @@ pytest-asyncio
|
||||||
pymupdf
|
pymupdf
|
||||||
python-docx
|
python-docx
|
||||||
scikit-learn
|
scikit-learn
|
||||||
aiofiles
|
aiofiles
|
||||||
|
paddlepaddle
|
||||||
|
paddleocr
|
||||||
|
pillow
|
||||||
|
|
@ -1,25 +1,55 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import docx
|
import docx
|
||||||
import fitz # PyMuPDF
|
import fitz # PyMuPDF
|
||||||
from pathlib import Path
|
|
||||||
import aiofiles
|
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:
|
async with aiofiles.open(path, "r", encoding="utf-8") as f:
|
||||||
return await f.read()
|
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:
|
async with aiofiles.open(path, "r", encoding="utf-8") as f:
|
||||||
return await f.read()
|
return await f.read()
|
||||||
|
|
||||||
|
|
||||||
async def load_docx(path: str) -> str:
|
async def _load_docx(path: str) -> str:
|
||||||
# docx 读取是阻塞 I/O → 放入线程池
|
# docx 读取是阻塞 I/O → 放入线程池
|
||||||
def _read():
|
def _read():
|
||||||
doc = docx.Document(path)
|
doc = docx.Document(path)
|
||||||
|
|
@ -28,29 +58,56 @@ async def load_docx(path: str) -> str:
|
||||||
return await asyncio.to_thread(_read)
|
return await asyncio.to_thread(_read)
|
||||||
|
|
||||||
|
|
||||||
async def load_pdf(path: str) -> str:
|
async def _load_pdf(path: str) -> str:
|
||||||
# PyMuPDF 也是阻塞 → 放入线程池
|
|
||||||
def _read():
|
|
||||||
pdf = fitz.open(path)
|
|
||||||
return "\n".join(page.get_text() for page in pdf)
|
|
||||||
|
|
||||||
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:
|
async def load_document(path: str) -> str:
|
||||||
suffix = Path(path).suffix.lower()
|
suffix = Path(path).suffix.lower()
|
||||||
|
|
||||||
|
print(f"读取文件: {path}")
|
||||||
|
|
||||||
if suffix == ".txt":
|
if suffix == ".txt":
|
||||||
return await load_txt(path)
|
return await _load_txt(path)
|
||||||
if suffix in (".md", ".markdown"):
|
if suffix in (".md", ".markdown"):
|
||||||
return await load_md(path)
|
return await _load_md(path)
|
||||||
if suffix == ".docx":
|
if suffix == ".docx":
|
||||||
return await load_docx(path)
|
return await _load_docx(path)
|
||||||
if suffix == ".pdf":
|
if suffix == ".pdf":
|
||||||
return await load_pdf(path)
|
return await _load_pdf(path)
|
||||||
|
|
||||||
raise ValueError(f"Unsupported file type: {suffix}")
|
raise ValueError(f"Unsupported file type: {suffix}")
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,7 @@ async def test_embedding():
|
||||||
|
|
||||||
shared = {
|
shared = {
|
||||||
"files": [
|
"files": [
|
||||||
|
"./files/带图片的.pdf",
|
||||||
"./files/大白智问-API接入文档-V1.2.2.pdf",
|
"./files/大白智问-API接入文档-V1.2.2.pdf",
|
||||||
"./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md",
|
"./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md",
|
||||||
"./files/deepsearch状态.txt",
|
"./files/deepsearch状态.txt",
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue