240 lines
11 KiB
Python
240 lines
11 KiB
Python
"""Offline, request-private OCR worker. Paths and text arrive only on stdin.
|
|
|
|
Never fetch a model or URL. The PHP owner enforces the hard process deadline;
|
|
checkpoints keep already completed sources available when native inference stops.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import importlib.metadata
|
|
import importlib.util
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import socket
|
|
import stat
|
|
import sys
|
|
import time
|
|
import warnings
|
|
from pathlib import Path
|
|
|
|
MAX_BYTES = 20 * 1024 * 1024
|
|
MAX_PIXELS = 40_000_000
|
|
MAX_OUTPUT = 512 * 1024
|
|
MAX_INPUT = 1024 * 1024
|
|
MODEL_NAMES = {
|
|
"Det": "PP-OCRv6_det_small.onnx",
|
|
"Cls": "ch_ppocr_mobile_v2.0_cls_mobile.onnx",
|
|
"Rec": "PP-OCRv6_rec_small.onnx",
|
|
}
|
|
|
|
|
|
class Unavailable(Exception):
|
|
"""A fixed, non-sensitive failure code."""
|
|
|
|
|
|
def deny_network(*_args, **_kwargs):
|
|
raise OSError("network_disabled")
|
|
|
|
|
|
def validated_request(request):
|
|
if not isinstance(request, dict):
|
|
raise Unavailable("invalid_request")
|
|
root_value = request.get("allowed_root")
|
|
if not isinstance(root_value, str) or not root_value or "\0" in root_value or root_value.startswith(("//", "\\\\")) or not Path(root_value).is_absolute():
|
|
raise Unavailable("invalid_request")
|
|
root = Path(root_value).resolve(strict=True)
|
|
images = request.get("images")
|
|
if not root.is_dir() or not isinstance(images, list) or len(images) > 256:
|
|
raise Unavailable("invalid_request")
|
|
seen, rows = set(), []
|
|
for item in images:
|
|
if not isinstance(item, dict) or type(item.get("index")) is not int or not 0 < item["index"] <= 1_000_000_000:
|
|
raise Unavailable("invalid_request")
|
|
if item["index"] in seen:
|
|
raise Unavailable("invalid_request")
|
|
seen.add(item["index"])
|
|
value = item.get("path")
|
|
if not isinstance(value, str) or not value or "\0" in value or "://" in value or value.startswith(("//", "\\\\")) or not Path(value).is_absolute():
|
|
raise Unavailable("invalid_source")
|
|
path = Path(value).resolve(strict=True)
|
|
metadata = path.stat()
|
|
if root not in path.parents or not stat.S_ISREG(metadata.st_mode) or not 0 < metadata.st_size <= MAX_BYTES:
|
|
raise Unavailable("invalid_source")
|
|
rows.append((item["index"], path))
|
|
threads, detection = request.get("threads", 1), request.get("detection_limit", 640)
|
|
duration = request.get("deadline_ms", 0)
|
|
budget = request.get("max_text_bytes", 131072)
|
|
if type(threads) is not int or threads not in (1, 2) or type(detection) is not int or detection not in (640, 736):
|
|
raise Unavailable("invalid_request")
|
|
if type(duration) is not int or not 0 <= duration <= 300000 or type(budget) is not int or not 1 <= budget <= 262144:
|
|
raise Unavailable("invalid_request")
|
|
return rows, threads, detection, duration / 1000, budget
|
|
|
|
|
|
def load_image(path):
|
|
from PIL import Image, ImageOps
|
|
|
|
Image.MAX_IMAGE_PIXELS = MAX_PIXELS
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error", Image.DecompressionBombWarning)
|
|
with Image.open(path) as image:
|
|
if not 0 < image.width * image.height <= MAX_PIXELS:
|
|
raise Unavailable("image_limits")
|
|
if max(image.size) > 12000 or max(image.size) / min(image.size) > 32:
|
|
raise Unavailable("image_limits")
|
|
if image.format not in {"PNG", "JPEG", "WEBP", "BMP", "TIFF", "GIF"}:
|
|
raise Unavailable("invalid_source")
|
|
if getattr(image, "n_frames", 1) != 1:
|
|
raise Unavailable("multi_frame_unsupported")
|
|
image.seek(0)
|
|
return ImageOps.exif_transpose(image).convert("RGB")
|
|
|
|
|
|
def load_engine(threads, detection):
|
|
for package, version in (("rapidocr", "3.10.0"), ("onnxruntime", "1.31.0")):
|
|
if importlib.metadata.version(package) != version:
|
|
raise Unavailable("dependency_unavailable")
|
|
spec = importlib.util.find_spec("rapidocr")
|
|
if spec is None or not spec.origin:
|
|
raise Unavailable("dependency_unavailable")
|
|
models = Path(spec.origin).resolve().parent / "models"
|
|
params = {
|
|
"Global.log_level": "critical", "Det.limit_side_len": detection, "Det.limit_type": "max",
|
|
"EngineConfig.onnxruntime.intra_op_num_threads": threads,
|
|
"EngineConfig.onnxruntime.inter_op_num_threads": 1,
|
|
"EngineConfig.onnxruntime.use_cuda": False, "EngineConfig.onnxruntime.use_dml": False,
|
|
"EngineConfig.onnxruntime.use_cann": False, "EngineConfig.onnxruntime.use_coreml": False,
|
|
}
|
|
for stage, name in MODEL_NAMES.items():
|
|
path = (models / name).resolve(strict=True)
|
|
if models.resolve() not in path.parents or not path.is_file():
|
|
raise Unavailable("model_unavailable")
|
|
params[f"{stage}.model_path"] = str(path)
|
|
import cv2
|
|
from rapidocr import RapidOCR
|
|
|
|
cv2.setNumThreads(threads)
|
|
return RapidOCR(params=params)
|
|
|
|
|
|
def normalized_lines(detected, budget):
|
|
texts, scores, boxes = detected.txts, detected.scores, detected.boxes
|
|
if texts is None or scores is None or boxes is None or not len(texts):
|
|
raise Unavailable("no_text")
|
|
if not len(texts) == len(scores) == len(boxes) or len(texts) > 1024:
|
|
raise Unavailable("invalid_output")
|
|
lines, used = [], 0
|
|
for text, score, raw_box in zip(texts, scores, boxes, strict=True):
|
|
if not isinstance(text, str) or not text.strip() or any(ord(char) < 32 and char not in "\n\t" for char in text):
|
|
raise Unavailable("invalid_output")
|
|
size = len(text.encode("utf-8"))
|
|
if size > 8192 or used + size > budget:
|
|
raise Unavailable("text_limit")
|
|
confidence = float(score)
|
|
box = raw_box.tolist() if hasattr(raw_box, "tolist") else raw_box
|
|
if not math.isfinite(confidence) or not 0 <= confidence <= 1 or not isinstance(box, list) or len(box) != 4:
|
|
raise Unavailable("invalid_output")
|
|
safe_box = []
|
|
for point in box:
|
|
if not isinstance(point, (list, tuple)) or len(point) != 2:
|
|
raise Unavailable("invalid_output")
|
|
coordinates = [float(value) for value in point]
|
|
if any(not math.isfinite(value) or abs(value) > 1_000_000 for value in coordinates):
|
|
raise Unavailable("invalid_output")
|
|
safe_box.append(coordinates)
|
|
lines.append({"text": text, "confidence": confidence, "box": safe_box})
|
|
used += size
|
|
return lines, used
|
|
|
|
|
|
def run(request, *, engine_factory=load_engine, image_loader=load_image, checkpoint=None):
|
|
started = time.monotonic()
|
|
try:
|
|
sources, threads, detection, duration, text_budget = validated_request(request)
|
|
except Exception: # noqa: BLE001 - The process boundary must never expose private paths.
|
|
return {"status": "unavailable", "reason": "invalid_request", "images": []}
|
|
deadline = started + duration
|
|
rows = [{"index": index, "status": "unreadable", "reason": "not_processed", "lines": []}
|
|
for index, _ in sources]
|
|
result = {"status": "partial", "reason": "not_processed", "images": rows}
|
|
if not sources:
|
|
return {"status": "completed", "images": []}
|
|
if duration <= 0:
|
|
return {**result, "status": "unavailable", "reason": "deadline"}
|
|
for name in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS"):
|
|
os.environ[name] = str(threads)
|
|
logging.disable(logging.CRITICAL)
|
|
original = (socket.socket.connect, socket.socket.connect_ex, socket.create_connection, socket.getaddrinfo)
|
|
socket.socket.connect = socket.socket.connect_ex = socket.create_connection = socket.getaddrinfo = deny_network
|
|
try:
|
|
with open(os.devnull, "w", encoding="utf-8") as quiet, contextlib.redirect_stdout(quiet), contextlib.redirect_stderr(quiet):
|
|
try:
|
|
engine = engine_factory(threads, detection)
|
|
except Exception: # noqa: BLE001 - Optional native engines may raise arbitrary errors.
|
|
return {**result, "status": "unavailable", "reason": "dependency_unavailable"}
|
|
for position, (_, path) in enumerate(sources):
|
|
if time.monotonic() >= deadline:
|
|
result["reason"] = "deadline"
|
|
break
|
|
try:
|
|
with open(os.devnull, "w", encoding="utf-8") as quiet, contextlib.redirect_stdout(quiet), contextlib.redirect_stderr(quiet):
|
|
image = image_loader(path)
|
|
try:
|
|
import numpy as np
|
|
|
|
detected = engine(np.asarray(image))
|
|
finally:
|
|
image.close()
|
|
lines, used = normalized_lines(detected, text_budget)
|
|
candidate = {**rows[position], "status": "processed", "lines": lines}
|
|
candidate.pop("reason", None)
|
|
trial = {**result, "images": [candidate if index == position else row for index, row in enumerate(rows)]}
|
|
if len(json.dumps(trial, ensure_ascii=False, allow_nan=False).encode("utf-8")) > MAX_OUTPUT:
|
|
raise Unavailable("output_limit")
|
|
rows[position] = candidate
|
|
text_budget -= used
|
|
except Exception as error: # noqa: BLE001 - One decoder failure must not erase other sources.
|
|
reason = str(error) if isinstance(error, Unavailable) else "inference_failed"
|
|
rows[position] = {**rows[position], "reason": reason}
|
|
if checkpoint:
|
|
checkpoint(result)
|
|
processed = sum(row["status"] == "processed" for row in rows)
|
|
result["status"] = "completed" if processed == len(rows) else ("partial" if processed else "unavailable")
|
|
if result["status"] == "completed":
|
|
result.pop("reason", None)
|
|
elif result["reason"] != "deadline":
|
|
result["reason"] = "unreadable_sources"
|
|
return result
|
|
finally:
|
|
socket.socket.connect, socket.socket.connect_ex, socket.create_connection, socket.getaddrinfo = original
|
|
|
|
|
|
def main():
|
|
stream = sys.stdout
|
|
stream.reconfigure(encoding="utf-8", errors="strict")
|
|
|
|
def write(result):
|
|
data = json.dumps(result, ensure_ascii=False, allow_nan=False)
|
|
if stream.seekable():
|
|
stream.seek(0)
|
|
stream.truncate()
|
|
stream.write(data)
|
|
stream.flush()
|
|
|
|
try:
|
|
raw = sys.stdin.buffer.read(MAX_INPUT + 1)
|
|
if len(raw) > MAX_INPUT:
|
|
raise Unavailable("invalid_request")
|
|
request = json.loads(raw)
|
|
result = run(request, checkpoint=write if stream.seekable() else None)
|
|
except Exception: # noqa: BLE001 - Emit only the fixed safe wire error at the worker boundary.
|
|
result = {"status": "unavailable", "reason": "invalid_request", "images": []}
|
|
write(result)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|