mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix: make OCR portable and reliable across AMD64 and ARM64 (#519)
* fix: make OCR portable and reliable * fix: harden OCR installation portability * fix: pin OCR partials across downloads * fix: make OCR execution reliably asynchronous * fix: harden OCR portability and docs routes * fix: preserve decoder and docs safeguards
This commit is contained in:
SnapOtter
@@ -66,8 +66,6 @@ ALLOWED_SCRIPTS = {
|
||||
"inpaint",
|
||||
"install_feature",
|
||||
"noise_removal",
|
||||
"ocr",
|
||||
"ocr_pdf",
|
||||
"ocr_preprocess",
|
||||
"outpaint",
|
||||
"red_eye_removal",
|
||||
@@ -122,8 +120,6 @@ TOOL_BUNDLE_MAP = {
|
||||
"enhance_faces": "upscale-enhance",
|
||||
"noise_removal": "upscale-enhance",
|
||||
"restore": "photo-restoration",
|
||||
"ocr": "ocr",
|
||||
"ocr_pdf": "ocr",
|
||||
"transcribe": "transcription",
|
||||
}
|
||||
|
||||
|
||||
+11
-43
@@ -36,10 +36,10 @@ def gpu_available():
|
||||
"""Return True if a usable CUDA GPU is present at runtime.
|
||||
|
||||
This is the general "can any framework use a GPU" check (torch, then ONNX
|
||||
Runtime, then paddle). Tools bound to a single framework should instead call
|
||||
Runtime). Tools bound to a single framework should instead call
|
||||
the matching per-framework helper (torch_gpu_available,
|
||||
ctranslate2_gpu_available) so a GPU that only paddle or ONNX can use is not
|
||||
mistaken for a torch GPU.
|
||||
ctranslate2_gpu_available) so a GPU that only ONNX can use is not mistaken
|
||||
for a torch GPU.
|
||||
"""
|
||||
if _override_disables_gpu():
|
||||
return False
|
||||
@@ -55,16 +55,14 @@ def gpu_available():
|
||||
if onnx_available:
|
||||
return True
|
||||
|
||||
# A GPU is physically present but neither torch nor ONNX Runtime can use it.
|
||||
# The OCR bundle ships paddlepaddle-gpu, which still can, so probe paddle in
|
||||
# an isolated subprocess. This runs only now that nvidia-smi confirms a GPU,
|
||||
# and never in-process, because a GPU-less paddle import segfaults.
|
||||
# A GPU is physically present but neither supported framework can use it.
|
||||
# OCR is now an isolated portable CPU ONNX runtime, so Paddle must never be
|
||||
# imported as a fallback probe here (the old GPU wheel could crash while
|
||||
# resolving libcuda on otherwise valid CPU-only hosts).
|
||||
gpu_name = _nvidia_smi_gpu_name()
|
||||
if gpu_name:
|
||||
if _try_paddle_cuda_subprocess():
|
||||
return True
|
||||
print(f"[gpu] nvidia-smi found GPU ({gpu_name}) but neither torch, ONNX "
|
||||
"Runtime, nor paddle can use it; reinstall AI features for GPU support",
|
||||
print(f"[gpu] nvidia-smi found GPU ({gpu_name}) but neither torch nor ONNX "
|
||||
"Runtime can use it; reinstall the relevant AI feature for GPU support",
|
||||
file=sys.stderr, flush=True)
|
||||
return False
|
||||
|
||||
@@ -129,42 +127,12 @@ def _try_onnx_cuda():
|
||||
return False
|
||||
|
||||
|
||||
def _try_paddle_cuda_subprocess():
|
||||
"""Check GPU via paddle in an isolated subprocess.
|
||||
|
||||
The OCR bundle ships paddlepaddle-gpu with no torch or ONNX Runtime, so paddle
|
||||
is the only framework that can see the GPU on an OCR-only host. Importing
|
||||
paddlepaddle-gpu in-process segfaults on a GPU-less machine, so this runs in a
|
||||
throwaway subprocess and callers must confirm a GPU is present (via nvidia-smi)
|
||||
before invoking it. Returns True only when paddle has a CUDA build and a
|
||||
visible GPU. The result is signalled through the exit code so paddle's own
|
||||
import chatter on stdout cannot corrupt the reading.
|
||||
"""
|
||||
probe = (
|
||||
"import paddle, sys; "
|
||||
"sys.exit(0 if (paddle.is_compiled_with_cuda() "
|
||||
"and paddle.device.cuda.device_count() > 0) else 1)"
|
||||
)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", probe],
|
||||
capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return False
|
||||
if result.returncode == 0:
|
||||
print("[gpu] CUDA available via paddle (paddlepaddle-gpu)",
|
||||
file=sys.stderr, flush=True)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def torch_gpu_available():
|
||||
"""True iff torch itself can use CUDA (honors the SNAPOTTER_GPU override).
|
||||
|
||||
Torch-based tools (upscale, denoise, face enhancement, restore) must gate on
|
||||
this rather than gpu_available(), which can report True based on paddle or
|
||||
ONNX Runtime while torch is a CPU-only build. Routing those tools to CUDA on a
|
||||
this rather than gpu_available(), which can report True based on ONNX Runtime
|
||||
while torch is a CPU-only build. Routing those tools to CUDA on a
|
||||
device torch cannot use would crash them.
|
||||
"""
|
||||
if _override_disables_gpu():
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,435 +0,0 @@
|
||||
"""Text extraction from images using Tesseract, PaddleOCR PP-OCRv5, or PaddleOCR-VL 1.5."""
|
||||
import sys
|
||||
import json
|
||||
import os
|
||||
|
||||
# Prevent PaddlePaddle C++ runtime from probing for CUDA on CPU-only systems.
|
||||
# Without these, paddlepaddle-gpu can segfault during import on machines without
|
||||
# a GPU, because the C++ layer attempts GPU initialization before Python-level
|
||||
# device routing takes effect. Must run before any PaddleOCR import.
|
||||
from gpu import gpu_available
|
||||
if not gpu_available():
|
||||
if not os.environ.get("FLAGS_use_cuda"):
|
||||
os.environ["FLAGS_use_cuda"] = "0"
|
||||
if not os.environ.get("FLAGS_use_cudnn"):
|
||||
os.environ["FLAGS_use_cudnn"] = "0"
|
||||
|
||||
# Lazy-loaded VLM instance (stays resident in dispatcher process)
|
||||
_paddleocr_vl_instance = None
|
||||
|
||||
|
||||
def emit_progress(percent, stage):
|
||||
"""Emit structured progress to stderr for bridge.ts to capture."""
|
||||
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
# OCR quality tiers backed by PaddleOCR. PaddleOCR ships as the GPU build
|
||||
# (paddlepaddle-gpu) in the amd64 bundle; its native libs dlopen libcuda.so.1 at
|
||||
# import and segfault on a CPU-only host, so these tiers need a usable GPU.
|
||||
PADDLE_QUALITY_TIERS = ("balanced", "best")
|
||||
|
||||
|
||||
def effective_quality(requested):
|
||||
"""Return the OCR quality tier that can actually run on this host.
|
||||
|
||||
On a CPU-only host the PaddleOCR tiers (balanced/best) cannot load, so they
|
||||
transparently fall back to "fast" (Tesseract), which runs on CPU. This keeps a
|
||||
GPU-less host from importing paddlepaddle-gpu, whose import segfaults and wedges
|
||||
the shared AI dispatcher.
|
||||
"""
|
||||
if requested in PADDLE_QUALITY_TIERS and not gpu_available():
|
||||
return "fast"
|
||||
return requested
|
||||
|
||||
|
||||
TESSERACT_LANG_MAP = {
|
||||
"en": "eng", "de": "deu", "fr": "fra", "es": "spa",
|
||||
"zh": "chi_sim", "ja": "jpn", "ko": "kor",
|
||||
}
|
||||
|
||||
PADDLE_LANG_MAP = {
|
||||
"en": "en", "de": "latin", "fr": "latin", "es": "latin",
|
||||
"zh": "ch", "ja": "japan", "ko": "korean",
|
||||
}
|
||||
|
||||
# Bundled PaddleOCR models (shipped by the OCR feature bundle into MODELS_PATH).
|
||||
# Pinning the constructor at these dirs keeps OCR fully offline / air-gapped and
|
||||
# skips slow HuggingFace model resolution on first use. PP-OCRv5 server rec
|
||||
# covers Chinese+English; latin covers en/de/fr/es; plus a dedicated Korean rec.
|
||||
PADDLE_DET_MODEL = "PP-OCRv5_server_det"
|
||||
PADDLE_TEXTLINE_MODEL = "PP-LCNet_x1_0_textline_ori"
|
||||
PADDLE_REC_MODEL = {
|
||||
"ch": "PP-OCRv5_server_rec",
|
||||
"en": "latin_PP-OCRv5_mobile_rec",
|
||||
"latin": "latin_PP-OCRv5_mobile_rec",
|
||||
"korean": "korean_PP-OCRv5_mobile_rec",
|
||||
}
|
||||
|
||||
|
||||
def _bundled_paddle_kwargs(paddle_lang):
|
||||
"""Build PaddleOCR kwargs that use the bundled models in MODELS_PATH.
|
||||
|
||||
Only pins a component when its model is actually present on disk, so a
|
||||
partial bundle (or an unbundled language such as Japanese) falls back to
|
||||
PaddleOCR's default resolution for that component. The doc-orientation and
|
||||
doc-unwarping models are not bundled and not needed for plain OCR, so they
|
||||
are disabled to avoid a runtime HuggingFace download.
|
||||
"""
|
||||
models_dir = os.environ.get("MODELS_PATH", "/data/ai/models")
|
||||
|
||||
def model_dir(name):
|
||||
if not name:
|
||||
return None
|
||||
path = os.path.join(models_dir, name)
|
||||
return path if os.path.isdir(path) else None
|
||||
|
||||
kwargs = {"use_doc_orientation_classify": False, "use_doc_unwarping": False}
|
||||
|
||||
det = model_dir(PADDLE_DET_MODEL)
|
||||
if det:
|
||||
kwargs["text_detection_model_name"] = PADDLE_DET_MODEL
|
||||
kwargs["text_detection_model_dir"] = det
|
||||
|
||||
rec_name = PADDLE_REC_MODEL.get(paddle_lang)
|
||||
rec = model_dir(rec_name)
|
||||
if rec:
|
||||
kwargs["text_recognition_model_name"] = rec_name
|
||||
kwargs["text_recognition_model_dir"] = rec
|
||||
|
||||
textline = model_dir(PADDLE_TEXTLINE_MODEL)
|
||||
if textline:
|
||||
kwargs["textline_orientation_model_name"] = PADDLE_TEXTLINE_MODEL
|
||||
kwargs["textline_orientation_model_dir"] = textline
|
||||
kwargs["use_textline_orientation"] = True
|
||||
else:
|
||||
kwargs["use_textline_orientation"] = False
|
||||
|
||||
return kwargs
|
||||
|
||||
|
||||
def auto_detect_language(input_path):
|
||||
"""Detect the predominant script in the image using Tesseract multi-lang.
|
||||
|
||||
Runs a quick Tesseract pass with all installed language packs,
|
||||
then analyzes the Unicode character ranges in the output to
|
||||
determine which PaddleOCR language model to use.
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["tesseract", input_path, "stdout", "-l", "eng+kor+chi_sim+jpn"],
|
||||
capture_output=True, text=True, timeout=30,
|
||||
)
|
||||
text = result.stdout.strip()
|
||||
if not text:
|
||||
return "en"
|
||||
|
||||
hangul = sum(1 for c in text if "\uAC00" <= c <= "\uD7AF" or "\u1100" <= c <= "\u11FF")
|
||||
cjk = sum(1 for c in text if "\u4E00" <= c <= "\u9FFF")
|
||||
hiragana = sum(1 for c in text if "\u3040" <= c <= "\u309F")
|
||||
katakana = sum(1 for c in text if "\u30A0" <= c <= "\u30FF")
|
||||
latin = sum(1 for c in text if c.isascii() and c.isalpha())
|
||||
|
||||
total = hangul + cjk + hiragana + katakana + latin
|
||||
if total == 0:
|
||||
return "en"
|
||||
|
||||
if hangul / total > 0.3:
|
||||
return "ko"
|
||||
if (hiragana + katakana) / total > 0.2:
|
||||
return "ja"
|
||||
if cjk / total > 0.3:
|
||||
return "zh"
|
||||
return "en"
|
||||
except Exception:
|
||||
return "en"
|
||||
|
||||
|
||||
def run_tesseract(input_path, language, is_auto=False):
|
||||
"""Run Tesseract OCR (Fast tier)."""
|
||||
import subprocess
|
||||
|
||||
# When auto-detected, use all installed language packs for best coverage
|
||||
if is_auto:
|
||||
tess_lang = "eng+kor+chi_sim+jpn+deu+fra+spa"
|
||||
else:
|
||||
tess_lang = TESSERACT_LANG_MAP.get(language, "eng")
|
||||
|
||||
emit_progress(30, "Scanning")
|
||||
result = subprocess.run(
|
||||
["tesseract", input_path, "stdout", "-l", tess_lang],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
emit_progress(70, "Extracting text")
|
||||
text = result.stdout.strip()
|
||||
if result.returncode != 0 and not text:
|
||||
raise RuntimeError(result.stderr.strip() or "Tesseract failed")
|
||||
return text
|
||||
|
||||
|
||||
def _extract_ocr_texts(results):
|
||||
"""Extract text from PaddleOCR 3.x result objects.
|
||||
|
||||
Handles multiple result formats across PaddleOCR versions:
|
||||
- 3.4.x: OCRResult with .json["res"]["rec_texts"]
|
||||
- Earlier: result objects with .res dict containing "text" list
|
||||
"""
|
||||
text_parts = []
|
||||
for res in results:
|
||||
# PaddleOCR 3.4.x format: OCRResult with .json dict
|
||||
if hasattr(res, "json") and isinstance(res.json, dict):
|
||||
inner = res.json.get("res", {})
|
||||
rec_texts = inner.get("rec_texts", [])
|
||||
if rec_texts:
|
||||
text_parts.extend(rec_texts)
|
||||
continue
|
||||
# Older format: .res dict with "text" list
|
||||
if hasattr(res, "res") and isinstance(res.res, dict):
|
||||
text_parts.extend(res.res.get("text", []))
|
||||
return "\n".join(text_parts)
|
||||
|
||||
|
||||
def run_paddleocr_v5(input_path, language):
|
||||
"""Run PaddleOCR PP-OCRv5 server models (Balanced tier)."""
|
||||
# GPU-only: paddlepaddle-gpu segfaults at import on a CPU-only host (libcuda
|
||||
# absent). Refuse before importing so the caller falls back to Tesseract.
|
||||
if not gpu_available():
|
||||
raise ImportError(
|
||||
"PaddleOCR (paddlepaddle-gpu) requires a GPU; the amd64 bundle ships the "
|
||||
"GPU build, which cannot load on a CPU-only host. Use quality=fast (Tesseract)."
|
||||
)
|
||||
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
||||
|
||||
stdout_fd = os.dup(1)
|
||||
os.dup2(2, 1)
|
||||
|
||||
try:
|
||||
import logging
|
||||
from paddleocr import PaddleOCR
|
||||
|
||||
# Suppress PaddleOCR internal logging (replaces removed show_log param)
|
||||
for name in ("ppocr", "paddleocr", "paddle"):
|
||||
logging.getLogger(name).setLevel(logging.ERROR)
|
||||
|
||||
paddle_lang = PADDLE_LANG_MAP.get(language, "en")
|
||||
device = "gpu:0" if gpu_available() else "cpu"
|
||||
|
||||
emit_progress(20, "Loading")
|
||||
mk = _bundled_paddle_kwargs(paddle_lang)
|
||||
# When a bundled recognizer is pinned, the model selects the script, so
|
||||
# we omit lang (this is the proven fully-offline path). Lang-based
|
||||
# resolution for a language without a bundled recognizer (e.g. ja) and
|
||||
# a missing detection model both make PaddleOCR fetch models over the
|
||||
# network, which strict offline mode blocks with a clear error.
|
||||
if "text_recognition_model_dir" not in mk:
|
||||
from offline_guard import ensure_download_allowed
|
||||
ensure_download_allowed(f"PaddleOCR recognition model for language '{language}'")
|
||||
mk["lang"] = paddle_lang
|
||||
if "text_detection_model_dir" not in mk:
|
||||
from offline_guard import ensure_download_allowed
|
||||
ensure_download_allowed(f"PaddleOCR text detection model ({PADDLE_DET_MODEL})")
|
||||
ocr = PaddleOCR(
|
||||
device=device,
|
||||
ocr_version="PP-OCRv5",
|
||||
enable_mkldnn=False,
|
||||
**mk,
|
||||
)
|
||||
emit_progress(30, "Scanning")
|
||||
results = ocr.predict(input=input_path)
|
||||
emit_progress(70, "Extracting text")
|
||||
|
||||
text = _extract_ocr_texts(results)
|
||||
finally:
|
||||
os.dup2(stdout_fd, 1)
|
||||
os.close(stdout_fd)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def run_paddleocr_vl(input_path):
|
||||
"""Run PaddleOCR-VL 1.5 vision-language model (Best tier).
|
||||
|
||||
The VLM is lazy-loaded on first call and stays resident in the
|
||||
dispatcher process for subsequent requests.
|
||||
Requires PaddlePaddle >= 3.2 for fused_rms_norm_ext.
|
||||
"""
|
||||
global _paddleocr_vl_instance
|
||||
# GPU-only: see run_paddleocr_v5. Refuse before importing paddle on CPU.
|
||||
if not gpu_available():
|
||||
raise ImportError(
|
||||
"PaddleOCR-VL (paddlepaddle-gpu) requires a GPU; the amd64 bundle ships the "
|
||||
"GPU build, which cannot load on a CPU-only host. Use quality=balanced or fast."
|
||||
)
|
||||
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
||||
|
||||
stdout_fd = os.dup(1)
|
||||
os.dup2(2, 1)
|
||||
|
||||
try:
|
||||
if _paddleocr_vl_instance is None:
|
||||
emit_progress(15, "Loading model")
|
||||
from paddleocr import PaddleOCRVL
|
||||
|
||||
device = "gpu" if gpu_available() else "cpu"
|
||||
_paddleocr_vl_instance = PaddleOCRVL(device=device)
|
||||
|
||||
emit_progress(30, "Scanning")
|
||||
output = _paddleocr_vl_instance.predict(input_path)
|
||||
emit_progress(70, "Extracting text")
|
||||
|
||||
text_parts = []
|
||||
for res in output:
|
||||
# PaddleOCR-VL 1.5+: markdown_texts holds the extracted text
|
||||
if hasattr(res, "markdown") and isinstance(res.markdown, dict):
|
||||
md_text = res.markdown.get("markdown_texts", "")
|
||||
if md_text:
|
||||
text_parts.append(md_text)
|
||||
continue
|
||||
if hasattr(res, "parsing_res_list"):
|
||||
for block in res.parsing_res_list:
|
||||
content = block.get("block_content", "")
|
||||
if content:
|
||||
text_parts.append(content)
|
||||
elif hasattr(res, "rec_text"):
|
||||
text_parts.append(res.rec_text)
|
||||
# Also try the json-based extraction as fallback
|
||||
elif hasattr(res, "json") and isinstance(res.json, dict):
|
||||
inner = res.json.get("res", {})
|
||||
rec_texts = inner.get("rec_texts", [])
|
||||
text_parts.extend(rec_texts)
|
||||
|
||||
text = "\n".join(text_parts)
|
||||
finally:
|
||||
os.dup2(stdout_fd, 1)
|
||||
os.close(stdout_fd)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def main():
|
||||
input_path = sys.argv[1]
|
||||
settings = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}
|
||||
|
||||
quality = settings.get("quality", None)
|
||||
language = settings.get("language", "auto")
|
||||
enhance = settings.get("enhance", True)
|
||||
|
||||
# Backward compat: old "engine" param maps to quality
|
||||
if quality is None:
|
||||
engine = settings.get("engine", "tesseract")
|
||||
quality = "fast" if engine == "tesseract" else "balanced"
|
||||
|
||||
# On a CPU-only host, downgrade GPU-only tiers (PaddleOCR) to Tesseract so we
|
||||
# never import paddlepaddle-gpu, whose import segfaults and wedges the dispatcher.
|
||||
downgraded = effective_quality(quality)
|
||||
if downgraded != quality:
|
||||
print(
|
||||
json.dumps({"info": f"{quality} OCR needs a GPU; using {downgraded} (Tesseract) on CPU"}),
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
quality = downgraded
|
||||
|
||||
preprocessed_path = None
|
||||
try:
|
||||
emit_progress(5, "Preparing")
|
||||
|
||||
# Preprocessing (if enabled)
|
||||
if enhance:
|
||||
emit_progress(8, "Enhancing image")
|
||||
try:
|
||||
from ocr_preprocess import preprocess
|
||||
preprocessed_path = input_path + "_enhanced.png"
|
||||
preprocess(input_path, preprocessed_path)
|
||||
input_path = preprocessed_path
|
||||
except Exception as e:
|
||||
print(json.dumps({"warning": f"Enhancement skipped: {e}"}), file=sys.stderr, flush=True)
|
||||
preprocessed_path = None
|
||||
|
||||
# Language auto-detection
|
||||
was_auto = language == "auto"
|
||||
if was_auto:
|
||||
emit_progress(10, "Detecting language")
|
||||
language = auto_detect_language(input_path)
|
||||
|
||||
engine_used = quality
|
||||
|
||||
# Route to engine based on quality tier
|
||||
if quality == "fast":
|
||||
try:
|
||||
text = run_tesseract(input_path, language, is_auto=was_auto)
|
||||
engine_used = "tesseract"
|
||||
except FileNotFoundError:
|
||||
print(json.dumps({"success": False, "error": "Tesseract is not installed"}))
|
||||
sys.exit(1)
|
||||
|
||||
elif quality == "balanced":
|
||||
try:
|
||||
text = run_paddleocr_v5(input_path, language)
|
||||
engine_used = "paddleocr-v5"
|
||||
except ImportError as e:
|
||||
print(json.dumps({
|
||||
"success": False,
|
||||
"error": (
|
||||
f"PaddleOCR is not installed: {e}. "
|
||||
"Install the OCR feature or use quality=fast for Tesseract."
|
||||
),
|
||||
}))
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(json.dumps({
|
||||
"success": False,
|
||||
"error": (
|
||||
f"PaddleOCR PP-OCRv5 failed: {type(e).__name__}: {e}. "
|
||||
"Install the OCR feature or use quality=fast for Tesseract."
|
||||
),
|
||||
}))
|
||||
sys.exit(1)
|
||||
|
||||
elif quality == "best":
|
||||
try:
|
||||
text = run_paddleocr_vl(input_path)
|
||||
engine_used = "paddleocr-vl"
|
||||
except ImportError as e:
|
||||
print(json.dumps({
|
||||
"success": False,
|
||||
"error": (
|
||||
f"PaddleOCR-VL is not available: {e}. "
|
||||
"Install the OCR feature or use quality=balanced for PP-OCRv5."
|
||||
),
|
||||
}))
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(json.dumps({
|
||||
"success": False,
|
||||
"error": (
|
||||
f"PaddleOCR-VL failed: {type(e).__name__}: {e}. "
|
||||
"Install the OCR feature or use quality=balanced for PP-OCRv5."
|
||||
),
|
||||
}))
|
||||
sys.exit(1)
|
||||
|
||||
else:
|
||||
print(json.dumps({"success": False, "error": f"Unknown quality: {quality}"}))
|
||||
sys.exit(1)
|
||||
|
||||
emit_progress(95, "Done")
|
||||
print(json.dumps({"success": True, "text": text, "engine": engine_used}))
|
||||
|
||||
except Exception as e:
|
||||
print(json.dumps({"success": False, "error": str(e)}))
|
||||
sys.exit(1)
|
||||
finally:
|
||||
# Clean up preprocessed temp file
|
||||
if preprocessed_path:
|
||||
try:
|
||||
os.remove(preprocessed_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,196 +0,0 @@
|
||||
"""OCR text extraction from PDF documents.
|
||||
|
||||
Rasterizes each page via PyMuPDF (fitz) and runs the same OCR engine
|
||||
pipeline as ocr.py (tesseract / PaddleOCR PP-OCRv5 / PaddleOCR-VL).
|
||||
|
||||
Contract: argv[1] = pdf path, argv[2] = options JSON {"quality","language","pages"}.
|
||||
Output: {"success": true, "engine": ..., "pages": n, "text": ...} to stdout.
|
||||
Errors: {"error": ...} to stdout, exit 1.
|
||||
|
||||
Engine functions are imported from sibling ocr.py (the dispatcher adds the
|
||||
scripts directory to sys.path, and `import ocr` loads ocr.py as a module
|
||||
without triggering its __main__ guard).
|
||||
"""
|
||||
import sys
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
MAX_PAGES = 50
|
||||
|
||||
|
||||
def emit_progress(percent, stage):
|
||||
"""Emit structured progress to stderr for bridge.ts to capture."""
|
||||
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
def parse_page_spec(spec, total_pages):
|
||||
"""Parse a page spec like 'all', '1-3,5', '2,4-6' into a sorted list of 0-based indices.
|
||||
|
||||
Returns (pages, error) where error is a string if the spec is invalid.
|
||||
"""
|
||||
if spec.strip().lower() == "all":
|
||||
return list(range(total_pages)), None
|
||||
|
||||
pages = set()
|
||||
parts = spec.split(",")
|
||||
for part in parts:
|
||||
part = part.strip()
|
||||
if not part:
|
||||
continue
|
||||
if "-" in part:
|
||||
bounds = part.split("-", 1)
|
||||
try:
|
||||
start = int(bounds[0].strip())
|
||||
end = int(bounds[1].strip())
|
||||
except ValueError:
|
||||
return None, f"Invalid page range: {part}"
|
||||
if start < 1 or end < 1:
|
||||
return None, f"Invalid page range: {part} (pages start at 1)"
|
||||
if start > end:
|
||||
return None, f"Invalid page range: {part} (start > end)"
|
||||
if end > total_pages:
|
||||
return None, f"Invalid page range: {part} (document has {total_pages} pages)"
|
||||
for i in range(start, end + 1):
|
||||
pages.add(i - 1) # 0-based
|
||||
else:
|
||||
try:
|
||||
num = int(part)
|
||||
except ValueError:
|
||||
return None, f"Invalid page number: {part}"
|
||||
if num < 1:
|
||||
return None, f"Invalid page number: {num} (pages start at 1)"
|
||||
if num > total_pages:
|
||||
return None, f"Invalid page number: {num} (document has {total_pages} pages)"
|
||||
pages.add(num - 1) # 0-based
|
||||
|
||||
if not pages:
|
||||
return None, "No pages specified"
|
||||
|
||||
return sorted(pages), None
|
||||
|
||||
|
||||
def main():
|
||||
input_path = sys.argv[1]
|
||||
settings = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}
|
||||
|
||||
quality = settings.get("quality", "balanced")
|
||||
language = settings.get("language", "auto")
|
||||
pages_spec = settings.get("pages", "all")
|
||||
|
||||
try:
|
||||
emit_progress(2, "Opening PDF")
|
||||
|
||||
# Lazy import: fitz (PyMuPDF) is in the base venv but keep it lazy
|
||||
# so py_compile succeeds without it.
|
||||
import fitz
|
||||
|
||||
doc = fitz.open(input_path)
|
||||
total_pages = len(doc)
|
||||
|
||||
if total_pages == 0:
|
||||
print(json.dumps({"error": "PDF has no pages"}))
|
||||
sys.exit(1)
|
||||
|
||||
# Parse page specification
|
||||
page_indices, parse_error = parse_page_spec(pages_spec, total_pages)
|
||||
if parse_error:
|
||||
print(json.dumps({"error": parse_error}))
|
||||
sys.exit(1)
|
||||
|
||||
if len(page_indices) > MAX_PAGES:
|
||||
print(json.dumps({"error": f"Too many pages for OCR (max {MAX_PAGES})"}))
|
||||
sys.exit(1)
|
||||
|
||||
emit_progress(5, "Preparing")
|
||||
|
||||
# Import engine functions from sibling ocr.py.
|
||||
# The dispatcher's scripts directory is on sys.path, so this import
|
||||
# loads ocr.py as a module (its __main__ guard prevents re-execution).
|
||||
import ocr as ocr_module
|
||||
|
||||
# Language auto-detection on the first page's raster
|
||||
was_auto = language == "auto"
|
||||
|
||||
# Create a temp directory for scratch PNGs next to the input
|
||||
scratch_dir = tempfile.mkdtemp(
|
||||
prefix="ocr_pdf_",
|
||||
dir=os.path.dirname(input_path) or tempfile.gettempdir(),
|
||||
)
|
||||
|
||||
page_texts = []
|
||||
engine_used = quality
|
||||
num_pages = len(page_indices)
|
||||
|
||||
for i, page_idx in enumerate(page_indices):
|
||||
page_num = page_idx + 1 # 1-based for display
|
||||
pct = 5 + int((i / num_pages) * 85)
|
||||
emit_progress(pct, f"Page {page_num}/{total_pages}")
|
||||
|
||||
# Rasterize page to PNG at 200 DPI
|
||||
page = doc[page_idx]
|
||||
pix = page.get_pixmap(dpi=200)
|
||||
png_path = os.path.join(scratch_dir, f"page_{page_idx}.png")
|
||||
pix.save(png_path)
|
||||
|
||||
# Auto-detect language on the first page only
|
||||
current_lang = language
|
||||
if was_auto and i == 0:
|
||||
current_lang = ocr_module.auto_detect_language(png_path)
|
||||
language = current_lang # reuse for remaining pages
|
||||
elif was_auto:
|
||||
current_lang = language # already detected
|
||||
|
||||
# Always route PDF OCR through tesseract. PaddleOCR segfaults
|
||||
# on arm64 CPU when processing rasterised PDF pages (SIGSEGV in
|
||||
# the doc-orientation / structural-analysis stage). Tesseract is
|
||||
# reliable for page-level images and is already installed in the
|
||||
# container with multi-language packs. The image-OCR tool still
|
||||
# offers PaddleOCR tiers for single images where it is stable.
|
||||
if quality not in ("fast", "balanced", "best"):
|
||||
print(json.dumps({"error": f"Unknown quality: {quality}"}))
|
||||
sys.exit(1)
|
||||
|
||||
text = ocr_module.run_tesseract(png_path, current_lang, is_auto=was_auto)
|
||||
engine_used = "tesseract"
|
||||
|
||||
page_texts.append((page_num, text))
|
||||
|
||||
# Clean up scratch PNG immediately
|
||||
try:
|
||||
os.remove(png_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
# Clean up scratch directory
|
||||
try:
|
||||
os.rmdir(scratch_dir)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
doc.close()
|
||||
|
||||
# Join page texts with page headers
|
||||
parts = []
|
||||
for page_num, text in page_texts:
|
||||
parts.append(f"\n\n--- Page {page_num} ---\n\n{text}")
|
||||
|
||||
# Strip leading newlines from the first page
|
||||
full_text = "".join(parts).strip()
|
||||
|
||||
emit_progress(95, "Done")
|
||||
print(json.dumps({
|
||||
"success": True,
|
||||
"engine": engine_used,
|
||||
"pages": num_pages,
|
||||
"text": full_text,
|
||||
}))
|
||||
|
||||
except Exception as e:
|
||||
print(json.dumps({"error": str(e)}))
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,263 @@
|
||||
"""Persistent JSON-line entrypoint for the isolated OCR runtime."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, BinaryIO, Callable, Mapping, TextIO
|
||||
|
||||
try:
|
||||
# Direct execution inside the immutable runtime generation.
|
||||
from ocr_runtime import OcrRuntime
|
||||
except ModuleNotFoundError: # pragma: no cover - exercised by package imports in tooling
|
||||
from .ocr_runtime import OcrRuntime
|
||||
|
||||
|
||||
PROTOCOL_VERSION = 1
|
||||
MAX_REQUEST_BYTES = 64 * 1024
|
||||
MAX_ERROR_CHARS = 400
|
||||
|
||||
|
||||
def configure_offline_environment() -> None:
|
||||
"""Set offline switches before importing any optional model libraries."""
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||||
os.environ["SNAPOTTER_OCR_OFFLINE"] = "1"
|
||||
os.environ["NO_PROXY"] = "*"
|
||||
os.environ["no_proxy"] = "*"
|
||||
|
||||
|
||||
def _request_id(request: Any) -> str:
|
||||
if isinstance(request, Mapping):
|
||||
value = request.get("requestId")
|
||||
if isinstance(value, str) and 0 < len(value) <= 128:
|
||||
return value
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _failure(request_id: str, code: str, message: str) -> dict[str, Any]:
|
||||
bounded = " ".join(str(message).split())[:MAX_ERROR_CHARS] or "OCR runtime request failed"
|
||||
return {
|
||||
"protocolVersion": PROTOCOL_VERSION,
|
||||
"requestId": request_id,
|
||||
"ok": False,
|
||||
"error": {"code": code, "message": bounded},
|
||||
}
|
||||
|
||||
|
||||
def _parse_settings(value: str) -> dict[str, Any]:
|
||||
try:
|
||||
settings = json.loads(value)
|
||||
except json.JSONDecodeError as error:
|
||||
raise ValueError("OCR settings are not valid JSON") from error
|
||||
if not isinstance(settings, dict):
|
||||
raise ValueError("OCR settings must be a JSON object")
|
||||
return settings
|
||||
|
||||
|
||||
def _parse_pages(value: str) -> tuple[tuple[int, Path], ...]:
|
||||
try:
|
||||
raw_pages = json.loads(value)
|
||||
except json.JSONDecodeError as error:
|
||||
raise ValueError("Rasterized pages are not valid JSON") from error
|
||||
if not isinstance(raw_pages, list) or not raw_pages:
|
||||
raise ValueError("Rasterized pages must be a non-empty JSON list")
|
||||
pages = []
|
||||
for index, item in enumerate(raw_pages):
|
||||
if isinstance(item, str):
|
||||
pages.append((index + 1, Path(item)))
|
||||
continue
|
||||
if (
|
||||
isinstance(item, dict)
|
||||
and isinstance(item.get("page"), int)
|
||||
and not isinstance(item.get("page"), bool)
|
||||
and isinstance(item.get("path"), str)
|
||||
):
|
||||
pages.append((item["page"], Path(item["path"])))
|
||||
continue
|
||||
raise ValueError("Each rasterized page must be a path or {page, path} object")
|
||||
return tuple(pages)
|
||||
|
||||
|
||||
def _validate_request(request: Any) -> tuple[str, str, list[str]]:
|
||||
if not isinstance(request, dict):
|
||||
raise ValueError("OCR runtime request must be an object")
|
||||
request_id = _request_id(request)
|
||||
if request_id == "unknown":
|
||||
raise ValueError("OCR runtime requestId is missing or invalid")
|
||||
protocol_version = request.get("protocolVersion")
|
||||
if type(protocol_version) is not int or protocol_version != PROTOCOL_VERSION:
|
||||
raise ValueError(f"Unsupported OCR runtime protocol version; expected {PROTOCOL_VERSION}")
|
||||
script = request.get("script")
|
||||
if script not in ("ocr", "ocr_pdf", "smoke"):
|
||||
raise ValueError("OCR runtime script must be ocr, ocr_pdf, or smoke")
|
||||
args = request.get("args")
|
||||
if not isinstance(args, list) or not all(isinstance(arg, str) for arg in args):
|
||||
raise ValueError("OCR runtime args must be a list of strings")
|
||||
if script == "smoke" and args:
|
||||
raise ValueError("OCR runtime smoke requires no arguments")
|
||||
if script != "smoke" and len(args) != 2:
|
||||
raise ValueError("OCR runtime scripts require exactly two arguments")
|
||||
return request_id, script, args
|
||||
|
||||
|
||||
def _prepare_request(request: Any) -> tuple[str, str, tuple[Any, ...]]:
|
||||
"""Validate and decode a frame without constructing model sessions."""
|
||||
request_id, script, args = _validate_request(request)
|
||||
if script == "smoke":
|
||||
return request_id, script, ()
|
||||
settings = _parse_settings(args[1])
|
||||
if script == "ocr":
|
||||
return request_id, script, (Path(args[0]), settings)
|
||||
return request_id, script, (_parse_pages(args[0]), settings)
|
||||
|
||||
|
||||
def _process_prepared_request(
|
||||
prepared: tuple[str, str, tuple[Any, ...]], runtime: OcrRuntime
|
||||
) -> dict[str, Any]:
|
||||
request_id, script, arguments = prepared
|
||||
try:
|
||||
if script == "smoke":
|
||||
result = runtime.smoke()
|
||||
elif script == "ocr":
|
||||
result = runtime.recognize_image(*arguments)
|
||||
else:
|
||||
result = runtime.recognize_pages(*arguments)
|
||||
return {
|
||||
"protocolVersion": PROTOCOL_VERSION,
|
||||
"requestId": request_id,
|
||||
"ok": True,
|
||||
"result": result,
|
||||
}
|
||||
except ValueError as error:
|
||||
return _failure(request_id, "invalid-request", str(error))
|
||||
except FileNotFoundError as error:
|
||||
return _failure(request_id, "file-not-found", str(error))
|
||||
except Exception as error:
|
||||
return _failure(request_id, "ocr-runtime-failed", str(error))
|
||||
|
||||
|
||||
def process_request(
|
||||
request: Any,
|
||||
*,
|
||||
runtime: OcrRuntime | None = None,
|
||||
runtime_factory: Callable[[], OcrRuntime] = OcrRuntime,
|
||||
) -> dict[str, Any]:
|
||||
request_id = _request_id(request)
|
||||
try:
|
||||
prepared = _prepare_request(request)
|
||||
request_id = prepared[0]
|
||||
active_runtime = runtime or runtime_factory()
|
||||
return _process_prepared_request(prepared, active_runtime)
|
||||
except ValueError as error:
|
||||
return _failure(request_id, "invalid-request", str(error))
|
||||
except FileNotFoundError as error:
|
||||
return _failure(request_id, "file-not-found", str(error))
|
||||
except Exception as error:
|
||||
return _failure(request_id, "ocr-runtime-failed", str(error))
|
||||
|
||||
|
||||
def main(
|
||||
argv: list[str] | None = None,
|
||||
*,
|
||||
runtime_factory: Callable[[], OcrRuntime] = OcrRuntime,
|
||||
stdin: BinaryIO | None = None,
|
||||
stdout: TextIO | None = None,
|
||||
stderr: TextIO | None = None,
|
||||
) -> int:
|
||||
configure_offline_environment()
|
||||
arguments = list(sys.argv[1:] if argv is None else argv)
|
||||
input_stream = sys.stdin.buffer if stdin is None else stdin
|
||||
output_stream = sys.stdout if stdout is None else stdout
|
||||
error_stream = sys.stderr if stderr is None else stderr
|
||||
|
||||
if arguments == ["--smoke"]:
|
||||
try:
|
||||
smoke_result = runtime_factory().smoke()
|
||||
except Exception as error:
|
||||
error_stream.write(f"OCR runtime smoke failed: {error}\n")
|
||||
error_stream.flush()
|
||||
return 1
|
||||
output_stream.write(
|
||||
json.dumps(
|
||||
{"smoke": True, **smoke_result},
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
output_stream.flush()
|
||||
return 0
|
||||
if arguments:
|
||||
error_stream.write("Usage: ocr_runner.py [--smoke]\n")
|
||||
error_stream.flush()
|
||||
return 2
|
||||
|
||||
runtime: OcrRuntime | None = None
|
||||
handled_requests = 0
|
||||
|
||||
while True:
|
||||
raw = input_stream.readline(MAX_REQUEST_BYTES + 1)
|
||||
if not raw:
|
||||
if handled_requests == 0:
|
||||
response = _failure(
|
||||
"unknown", "invalid-request", "OCR runtime request is empty"
|
||||
)
|
||||
output_stream.write(
|
||||
json.dumps(response, ensure_ascii=False, separators=(",", ":"))
|
||||
+ "\n"
|
||||
)
|
||||
output_stream.flush()
|
||||
return 1
|
||||
return 0
|
||||
|
||||
exit_after_response = False
|
||||
if len(raw) > MAX_REQUEST_BYTES:
|
||||
response = _failure(
|
||||
"unknown", "invalid-request", "OCR runtime request is too large"
|
||||
)
|
||||
# readline() leaves the rest of an oversized unterminated frame in
|
||||
# the stream. Exit after the bounded error instead of interpreting
|
||||
# the tail as another request and losing frame synchronization.
|
||||
exit_after_response = not raw.endswith(b"\n")
|
||||
else:
|
||||
try:
|
||||
request = json.loads(raw.decode("utf-8"))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as error:
|
||||
response = _failure(
|
||||
"unknown", "invalid-request", f"Malformed JSON request: {error}"
|
||||
)
|
||||
else:
|
||||
try:
|
||||
prepared = _prepare_request(request)
|
||||
except ValueError as error:
|
||||
response = _failure(
|
||||
_request_id(request), "invalid-request", str(error)
|
||||
)
|
||||
else:
|
||||
if runtime is None:
|
||||
try:
|
||||
runtime = runtime_factory()
|
||||
except Exception as error:
|
||||
response = _failure(
|
||||
_request_id(request), "ocr-runtime-failed", str(error)
|
||||
)
|
||||
exit_after_response = True
|
||||
else:
|
||||
response = _process_prepared_request(prepared, runtime)
|
||||
else:
|
||||
response = _process_prepared_request(prepared, runtime)
|
||||
|
||||
output_stream.write(
|
||||
json.dumps(response, ensure_ascii=False, separators=(",", ":")) + "\n"
|
||||
)
|
||||
output_stream.flush()
|
||||
handled_requests += 1
|
||||
if exit_after_response:
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -7,8 +7,6 @@
|
||||
# allowlisted model names (remove_bg.py ALLOWED_MODELS), never user paths.
|
||||
rembg==2.0.69
|
||||
realesrgan==0.3.0
|
||||
paddleocr==2.9.1
|
||||
paddlepaddle-gpu==3.0.0
|
||||
mediapipe>=0.10.21
|
||||
onnxruntime-gpu==1.20.1
|
||||
numpy==1.26.4
|
||||
|
||||
@@ -7,8 +7,6 @@
|
||||
# allowlisted model names (remove_bg.py ALLOWED_MODELS), never user paths.
|
||||
rembg[cpu]==2.0.69
|
||||
realesrgan==0.3.0
|
||||
paddleocr[doc-parser]>=3.4.0,<3.5.0
|
||||
paddlepaddle>=3.0.0,<3.1.0
|
||||
mediapipe>=0.10.21
|
||||
onnxruntime==1.20.1
|
||||
numpy==1.26.4
|
||||
|
||||
@@ -1,15 +1,5 @@
|
||||
"""gpu_available() must recognize a GPU that only paddle can use.
|
||||
|
||||
The OCR feature bundle ships paddlepaddle-gpu but no torch or ONNX Runtime. On an
|
||||
OCR-only GPU host the torch and ONNX probes both come up empty, so before this
|
||||
fix gpu_available() fell through to a plain nvidia-smi check that returned False
|
||||
by design, and OCR silently downgraded to Tesseract. gpu_available() now probes
|
||||
paddle too, in an isolated subprocess, but only after nvidia-smi confirms a GPU
|
||||
is physically present (importing paddlepaddle-gpu in-process segfaults on a
|
||||
GPU-less host and would wedge the shared AI dispatcher).
|
||||
"""
|
||||
"""Framework-specific CUDA probes must never depend on obsolete Paddle OCR."""
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
|
||||
@@ -17,32 +7,6 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
import gpu # noqa: E402
|
||||
|
||||
|
||||
# --- The isolated paddle probe --------------------------------------------
|
||||
|
||||
def test_paddle_probe_true_when_subprocess_reports_cuda(monkeypatch):
|
||||
def fake_run(cmd, **kwargs):
|
||||
return subprocess.CompletedProcess(cmd, returncode=0, stdout="", stderr="")
|
||||
|
||||
monkeypatch.setattr(gpu.subprocess, "run", fake_run)
|
||||
assert gpu._try_paddle_cuda_subprocess() is True
|
||||
|
||||
|
||||
def test_paddle_probe_false_when_subprocess_reports_no_cuda(monkeypatch):
|
||||
def fake_run(cmd, **kwargs):
|
||||
return subprocess.CompletedProcess(cmd, returncode=1, stdout="", stderr="")
|
||||
|
||||
monkeypatch.setattr(gpu.subprocess, "run", fake_run)
|
||||
assert gpu._try_paddle_cuda_subprocess() is False
|
||||
|
||||
|
||||
def test_paddle_probe_false_on_timeout(monkeypatch):
|
||||
def fake_run(cmd, **kwargs):
|
||||
raise subprocess.TimeoutExpired(cmd, 30)
|
||||
|
||||
monkeypatch.setattr(gpu.subprocess, "run", fake_run)
|
||||
assert gpu._try_paddle_cuda_subprocess() is False
|
||||
|
||||
|
||||
# --- gpu_available() orchestration -----------------------------------------
|
||||
|
||||
def _patch_probes(monkeypatch, torch, onnx, smi):
|
||||
@@ -54,39 +18,19 @@ def _patch_probes(monkeypatch, torch, onnx, smi):
|
||||
gpu.gpu_available.cache_clear()
|
||||
|
||||
|
||||
def test_gpu_available_true_when_only_paddle_sees_gpu(monkeypatch):
|
||||
# OCR-only GPU box: torch and ONNX absent, GPU present, paddle can use it.
|
||||
def test_gpu_available_does_not_treat_hardware_presence_as_framework_support(monkeypatch):
|
||||
_patch_probes(monkeypatch, torch=False, onnx=False, smi="NVIDIA GeForce RTX 4070")
|
||||
monkeypatch.setattr(gpu, "_try_paddle_cuda_subprocess", lambda: True)
|
||||
assert gpu.gpu_available() is True
|
||||
|
||||
|
||||
def test_gpu_available_false_when_paddle_cannot_use_gpu(monkeypatch):
|
||||
# GPU present but paddle is a CPU build or cannot init CUDA: stay conservative.
|
||||
_patch_probes(monkeypatch, torch=False, onnx=False, smi="NVIDIA GeForce RTX 4070")
|
||||
monkeypatch.setattr(gpu, "_try_paddle_cuda_subprocess", lambda: False)
|
||||
assert gpu.gpu_available() is False
|
||||
|
||||
|
||||
def test_gpu_available_never_probes_paddle_without_a_gpu(monkeypatch):
|
||||
# Safety invariant: on a GPU-less host a paddlepaddle-gpu import segfaults, so
|
||||
# the probe must never run when nvidia-smi finds no GPU.
|
||||
_patch_probes(monkeypatch, torch=False, onnx=False, smi=None)
|
||||
called = {"paddle": False}
|
||||
|
||||
def spy():
|
||||
called["paddle"] = True
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(gpu, "_try_paddle_cuda_subprocess", spy)
|
||||
assert gpu.gpu_available() is False
|
||||
assert called["paddle"] is False
|
||||
def test_gpu_module_has_no_paddle_probe():
|
||||
assert not hasattr(gpu, "_try_paddle_cuda_subprocess")
|
||||
|
||||
|
||||
# --- Per-framework detection (torch, ctranslate2) --------------------------
|
||||
#
|
||||
# Torch and CTranslate2 tools must gate on their OWN framework, not the general
|
||||
# gpu_available(), which can report True based on paddle or ONNX Runtime while
|
||||
# gpu_available(), which can report True based on ONNX Runtime while
|
||||
# torch is a CPU-only build. Consuming the shared boolean would make those tools
|
||||
# route to CUDA on a device their framework cannot use.
|
||||
|
||||
@@ -103,7 +47,7 @@ def test_torch_gpu_available_false_when_override_disables_gpu(monkeypatch):
|
||||
|
||||
|
||||
def test_torch_gpu_available_false_when_torch_is_cpu_only(monkeypatch):
|
||||
# The crux: gpu_available() may be True via paddle or ONNX on a GPU box, but a
|
||||
# The crux: gpu_available() may be True via ONNX on a GPU box, but a
|
||||
# CPU-only torch build must report no GPU so torch tools do not touch CUDA.
|
||||
monkeypatch.delenv("SNAPOTTER_GPU", raising=False)
|
||||
monkeypatch.setattr(gpu, "_try_torch_cuda", lambda: False)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,44 +0,0 @@
|
||||
"""OCR must degrade gracefully on CPU-only hosts.
|
||||
|
||||
The amd64 AI bundle ships paddlepaddle-gpu, whose native libraries dlopen
|
||||
libcuda.so.1 at import time and segfault on a host without a GPU (libcuda is the
|
||||
driver lib, injected only by nvidia-container-toolkit on GPU hosts). That segfault
|
||||
crashes the shared long-lived AI dispatcher and wedges all AI. So on a CPU-only
|
||||
host the PaddleOCR tiers (balanced/best) must transparently fall back to Tesseract
|
||||
and must never reach the paddle import.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
import ocr # noqa: E402
|
||||
|
||||
|
||||
def test_effective_quality_downgrades_paddle_tiers_on_cpu(monkeypatch):
|
||||
monkeypatch.setattr(ocr, "gpu_available", lambda: False)
|
||||
assert ocr.effective_quality("balanced") == "fast"
|
||||
assert ocr.effective_quality("best") == "fast"
|
||||
assert ocr.effective_quality("fast") == "fast"
|
||||
|
||||
|
||||
def test_effective_quality_preserves_paddle_tiers_on_gpu(monkeypatch):
|
||||
monkeypatch.setattr(ocr, "gpu_available", lambda: True)
|
||||
assert ocr.effective_quality("balanced") == "balanced"
|
||||
assert ocr.effective_quality("best") == "best"
|
||||
assert ocr.effective_quality("fast") == "fast"
|
||||
|
||||
|
||||
def test_run_paddleocr_v5_refuses_on_cpu_before_import(monkeypatch):
|
||||
# Must raise a GPU-specific error (the guard), NOT attempt the paddle import
|
||||
# that would segfault on a CPU-only host.
|
||||
monkeypatch.setattr(ocr, "gpu_available", lambda: False)
|
||||
with pytest.raises(ImportError, match="GPU"):
|
||||
ocr.run_paddleocr_v5("/nonexistent.png", "en")
|
||||
|
||||
|
||||
def test_run_paddleocr_vl_refuses_on_cpu_before_import(monkeypatch):
|
||||
monkeypatch.setattr(ocr, "gpu_available", lambda: False)
|
||||
with pytest.raises(ImportError, match="GPU"):
|
||||
ocr.run_paddleocr_vl("/nonexistent.png")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,18 @@
|
||||
"""Accurate OCR must never re-enter the mutable legacy Python dispatcher."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from packages.ai.python import dispatcher
|
||||
|
||||
|
||||
def test_legacy_dispatcher_does_not_allow_ocr_scripts():
|
||||
assert "ocr" not in dispatcher.ALLOWED_SCRIPTS
|
||||
assert "ocr_pdf" not in dispatcher.ALLOWED_SCRIPTS
|
||||
|
||||
|
||||
def test_shared_requirements_do_not_install_paddle():
|
||||
python_root = Path(__file__).resolve().parents[1]
|
||||
for name in ("requirements.txt", "requirements-gpu.txt"):
|
||||
requirements = (python_root / name).read_text(encoding="utf-8").lower()
|
||||
assert "paddleocr" not in requirements
|
||||
assert "paddlepaddle" not in requirements
|
||||
Reference in New Issue
Block a user