diff --git a/.env.example b/.env.example index a6a2881..df710a1 100644 --- a/.env.example +++ b/.env.example @@ -8,6 +8,7 @@ WISEGRADE_SECRET_KEY=replace-with-a-random-32-character-secret EMBEDDING_MODEL_NAME=BAAI/bge-m3 EMBEDDING_MODEL_PATH=/models/bge-m3 EMBEDDING_MODEL_VERSION=local +OCR_MODEL_DIR=/models/ocr MODEL_TIMEOUT_SECONDS=120 MODEL_CONCURRENCY=2 MODEL_RETRIES=2 diff --git a/backend/tests/test_parsing.py b/backend/tests/test_parsing.py index 8dd2195..07136a1 100644 --- a/backend/tests/test_parsing.py +++ b/backend/tests/test_parsing.py @@ -6,7 +6,7 @@ from docx import Document from pptx import Presentation from pptx.util import Inches -from wisegrade.parsing import parse_document +from wisegrade.parsing import _ocr_engine, parse_document def test_pdf_pages_and_images(tmp_path): @@ -20,6 +20,11 @@ def test_pdf_pages_and_images(tmp_path): assert (tmp_path / "assets" / "page-1.png").exists() +def test_scanned_pdf_requires_offline_ocr_package(tmp_path): + with pytest.raises(RuntimeError, match="离线OCR模型目录"): + _ocr_engine(str(tmp_path / "missing-ocr")) + + @pytest.mark.skipif(not shutil.which("libreoffice"), reason="LibreOffice is installed in the deployment image") def test_docx_and_pptx_render(tmp_path): word = Document() diff --git a/backend/wisegrade/config.py b/backend/wisegrade/config.py index 890a582..4d46579 100644 --- a/backend/wisegrade/config.py +++ b/backend/wisegrade/config.py @@ -19,6 +19,7 @@ class Settings(BaseSettings): embedding_model_name: str = "BAAI/bge-m3" embedding_model_path: Path = Path("/models/bge-m3") embedding_model_version: str = "local" + ocr_model_dir: Path = Path("/models/ocr") model_timeout_seconds: int = 120 model_concurrency: int = 2 model_retries: int = 2 diff --git a/backend/wisegrade/parsing.py b/backend/wisegrade/parsing.py index 6b6dfb1..dc7fb2b 100644 --- a/backend/wisegrade/parsing.py +++ b/backend/wisegrade/parsing.py @@ -1,11 +1,18 @@ import subprocess import tempfile +from functools import lru_cache from pathlib import Path +from threading import Lock import pymupdf as fitz from docx import Document from pptx import Presentation +from .config import get_settings + + +_ocr_lock = Lock() + def _render(page, output: Path): output.parent.mkdir(parents=True, exist_ok=True) @@ -13,15 +20,30 @@ def _render(page, output: Path): return str(output) -def _ocr(image_path: str): +@lru_cache(maxsize=1) +def _ocr_engine(model_dir: str): + root = Path(model_dir) + detection, recognition = root / "detection", root / "recognition" + if not detection.is_dir() or not recognition.is_dir(): + raise RuntimeError(f"扫描PDF需要离线OCR模型目录: {detection} 和 {recognition}") try: from paddleocr import PaddleOCR except ImportError as exc: raise RuntimeError("扫描PDF需要PaddleOCR;请安装OCR依赖或交由教师处理") from exc try: - engine = PaddleOCR(lang="ch", use_textline_orientation=True) + return PaddleOCR(lang="ch", text_detection_model_dir=str(detection), text_recognition_model_dir=str(recognition), + use_textline_orientation=False, use_doc_orientation_classify=False, use_doc_unwarping=False) except TypeError: - engine = PaddleOCR(use_angle_cls=True, lang="ch", show_log=False) + return PaddleOCR(lang="ch", det_model_dir=str(detection), rec_model_dir=str(recognition), use_angle_cls=False, show_log=False) + + +def _ocr(image_path: str): + engine = _ocr_engine(str(get_settings().ocr_model_dir)) + with _ocr_lock: + return _ocr_unlocked(engine, image_path) + + +def _ocr_unlocked(engine, image_path: str): if hasattr(engine, "predict"): result = engine.predict(image_path) lines = [] diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index ace0a85..a3177ec 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -59,7 +59,7 @@ EMBEDDING_MODEL_VERSION=your-package-version `EMBEDDING_MODEL_VERSION` 会写入每个评分批次快照。若模型目录缺失或内容不完整,导入仍可完成,知识库入库和评分会显示错误。可离线分发 `docker save` 生成的镜像包及单独的 `models-bge-m3.tar.zst`;目标机执行 `docker load -i docker-images.tar` 后解压模型包,运行相同的 Compose 配置。 -扫描 PDF 的 OCR 使用 PaddleOCR;需在离线环境准备其运行所需的 OCR 模型缓存。没有 OCR 模型且 PDF 无可提取文字时,文件进入解析失败状态并显示原因。 +扫描 PDF 的 OCR 使用 PaddleOCR。离线 OCR 包放在 `models/ocr/detection/` 和 `models/ocr/recognition/`,或修改 `OCR_MODEL_DIR` 指向包含这两个子目录的位置。没有 OCR 模型且 PDF 无可提取文字时,文件进入解析失败状态并显示原因;系统不会在解析时自动下载 OCR 模型。 ## 教师操作