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
@@ -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())
|
||||
Reference in New Issue
Block a user