"""Local-only, privacy-preserving comparison of classic OCR and PaddleOCR-VL.""" from __future__ import annotations import asyncio from io import BytesIO import json import os from pathlib import Path import shlex import subprocess import sys import tempfile import time import warnings import cv2 import numpy as np from PIL import Image, ImageOps, UnidentifiedImageError from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import FileResponse, JSONResponse from starlette.routing import Route from .fields import FIELD_CODES, assign_fields HERE = Path(__file__).resolve().parent PROJECT = HERE.parent MODEL_HOME = Path(os.environ.get("OCR_MODEL_HOME", PROJECT / ".cache")).expanduser().resolve() def _remote_command(seconds: int = 180, *, tiles_only: bool = False) -> str: remote_dir = os.environ["OCR_VISION_REMOTE_DIR"] remote_python = os.environ.get("OCR_VISION_REMOTE_PYTHON", "./.venv-vision/bin/python") runner = f'{shlex.quote(remote_python)} -m ocr_compare.vision_worker' if tiles_only: runner += ' --tiles-only' isolation = 'unshare -Urn ' if os.environ.get("OCR_VISION_REMOTE_UNSHARE") == "1" else "" device = shlex.quote(os.environ.get("OCR_VISION_REMOTE_DEVICE", "gpu:0")) model_home = os.environ.get("OCR_VISION_REMOTE_MODEL_HOME") model_env = f'OCR_MODEL_HOME={shlex.quote(model_home)} ' if model_home else "" return ( f'set -eu; cd {shlex.quote(remote_dir)}; directory=$(mktemp -d); ' 'trap \'rm -rf -- "$directory"\' EXIT; ' f'OCR_PRIVATE_DIR="$directory" OCR_VISION_DEVICE={device} {model_env}' f'timeout -k 5s {seconds}s {isolation}{runner}' ) MAX_UPLOAD = 16 * 1024 * 1024 MAX_PIXELS = 24_000_000 MAX_SIDE = 10_000 OCR_SIDE = 4000 Image.MAX_IMAGE_PIXELS = MAX_PIXELS _slot = asyncio.Semaphore(1) class InputError(ValueError): pass def _geometry(image: Image.Image, header: str | None) -> Image.Image: try: spec = json.loads(header) if header else {} rotation = int(spec.get("rotation", 0)) if rotation not in (0, 90, 180, 270): raise InputError("invalid_rotation") points = spec.get("corners") if points is not None: if (not isinstance(points, list) or len(points) != 4 or any(not isinstance(p, list) or len(p) != 2 or any(not isinstance(v, (int, float)) or isinstance(v, bool) or not 0 <= v <= 1 for v in p) for p in points)): raise InputError("invalid_corners") w, h = image.size source = np.array([[x * (w - 1), y * (h - 1)] for x, y in points], np.float32) area = cv2.contourArea(source) edges = [float(np.linalg.norm(source[(i + 1) % 4] - source[i])) for i in range(4)] turns = [] for i in range(4): ax, ay = source[(i + 1) % 4] - source[i] bx, by = source[(i + 2) % 4] - source[(i + 1) % 4] turns.append(float(ax * by - ay * bx)) if area < .02 * w * h or min(edges) < 30 or min(turns) <= 0: raise InputError("invalid_corners") out_w = round(max(edges[0], edges[2])) out_h = round(max(edges[1], edges[3])) if out_w < 30 or out_h < 30 or out_w * out_h > MAX_PIXELS: raise InputError("invalid_corners") # The four positions are always TL, TR, BR, BL. This avoids the # previous fixed-template aspect-ratio and rotation ambiguity. target = np.array([[0, 0], [out_w - 1, 0], [out_w - 1, out_h - 1], [0, out_h - 1]], np.float32) matrix = cv2.getPerspectiveTransform(source, target) warped = cv2.warpPerspective(np.asarray(image), matrix, (out_w, out_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(255, 255, 255)) image = Image.fromarray(warped) if rotation: image = image.rotate(-rotation, expand=True) if max(image.size) > OCR_SIDE: image.thumbnail((OCR_SIDE, OCR_SIDE), Image.Resampling.LANCZOS) return image except (ValueError, TypeError, KeyError, json.JSONDecodeError) as exc: if isinstance(exc, InputError): raise raise InputError("invalid_corners") from None def _prepare(data: bytes, header: str | None) -> tuple[bytes, tuple[int, int]]: try: with warnings.catch_warnings(): warnings.simplefilter("error", Image.DecompressionBombWarning) with Image.open(BytesIO(data)) as source: if source.format not in ("JPEG", "PNG", "MPO"): raise InputError("unsupported_encoding") w, h = source.size if min(w, h) < 1 or max(w, h) > MAX_SIDE or w * h > MAX_PIXELS: raise InputError("image_too_large") source.load() image = ImageOps.exif_transpose(source).convert("RGB") image = _geometry(image, header) output = BytesIO() image.save(output, format="JPEG", quality=92, subsampling=0) return output.getvalue(), image.size except InputError: raise except (OSError, ValueError, UnidentifiedImageError, Image.DecompressionBombError, Image.DecompressionBombWarning): raise InputError("invalid_image") from None def _classic(data: bytes, directory: Path, *, fine_text: bool = False) -> dict: started = time.monotonic() models = MODEL_HOME / "official_models" if any(not (models / name / "inference.pdiparams").is_file() for name in ("PP-OCRv5_mobile_det", "latin_PP-OCRv5_mobile_rec")): return {"status": "error", "error": "classic_models_missing", "elapsed_ms": 0} image_path = directory / "same-input.jpg" output_path = directory / "classic.json" image_path.write_bytes(data) env = os.environ.copy() env.update(PADDLE_PDX_CACHE_HOME=str(MODEL_HOME), HF_HOME=str(MODEL_HOME / "hf"), HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK="True", OMP_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1", OMP_THREAD_LIMIT="1") try: command = [sys.executable, "-m", "ocr_compare.classic_worker", str(image_path), str(output_path)] if fine_text: command.append("--fine-text") process = subprocess.run( command, cwd=PROJECT, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=120, check=False, ) except subprocess.TimeoutExpired: return {"status": "error", "error": "classic_timeout", "elapsed_ms": round((time.monotonic()-started)*1000)} except OSError: return {"status": "error", "error": "classic_failed", "elapsed_ms": round((time.monotonic()-started)*1000)} elapsed = round((time.monotonic() - started) * 1000) if process.returncode or not output_path.is_file(): return {"status": "error", "error": "classic_failed", "elapsed_ms": elapsed} try: raw = json.loads(output_path.read_text()) lines = [{"text": row["text"], "box": row["bbox"], "source": "OCR-Zeile"} for row in raw["lines"] if row.get("bbox")] except (OSError, ValueError, KeyError, TypeError): return {"status": "error", "error": "classic_invalid", "elapsed_ms": elapsed} model = "PP-OCRv5 mobile (Latin), lokal" + (" ยท Feintext" if fine_text else "") return {"status": "ok", "model": model, "elapsed_ms": elapsed, "lines": lines, "fields": assign_fields(lines), "regions": []} def _vision(data: bytes) -> dict: started = time.monotonic() mode = os.environ.get("OCR_VISION_MODE", "local").lower() if mode == "off": return {"status": "error", "error": "vision_disabled", "elapsed_ms": 0} if mode not in ("local", "ssh"): return {"status": "error", "error": "vision_config_error", "elapsed_ms": 0} if mode == "local" and any(not path.is_file() for path in ( MODEL_HOME / "official_models" / "PP-DocLayoutV3" / "inference.pdiparams", MODEL_HOME / "official_models" / "PaddleOCR-VL-1.6" / "model.safetensors")): return {"status": "error", "error": "vision_models_missing", "elapsed_ms": 0} if mode == "ssh" and (not os.environ.get("OCR_VISION_SSH_TARGET") or not os.environ.get("OCR_VISION_REMOTE_DIR")): return {"status": "error", "error": "vision_config_error", "elapsed_ms": 0} with Image.open(BytesIO(data)) as page: wide = page.width > 2500 and page.width / page.height > 1.65 limit = 90 if wide else 180 def run_worker(seconds: int, *, tiles_only: bool = False): if mode == "ssh": command = ["ssh", "-T", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", "-o", "ConnectTimeout=10"] identity = os.environ.get("OCR_VISION_SSH_IDENTITY") if identity: command += ["-i", str(Path(identity).expanduser()), "-o", "IdentitiesOnly=yes"] command += [os.environ["OCR_VISION_SSH_TARGET"], _remote_command(seconds, tiles_only=tiles_only)] env = None else: command = [os.environ.get("OCR_VISION_PYTHON", sys.executable), "-m", "ocr_compare.vision_worker"] if tiles_only: command.append("--tiles-only") env = os.environ.copy() env.update(OCR_MODEL_HOME=str(MODEL_HOME), OCR_VISION_DEVICE=os.environ.get("OCR_VISION_DEVICE", "cpu")) try: with tempfile.TemporaryDirectory(prefix="ocr-vision-") as private: if env is not None: env["OCR_PRIVATE_DIR"] = private return subprocess.run(command, input=data, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=seconds + 15, check=False, cwd=PROJECT if mode == "local" else None, env=env) except subprocess.TimeoutExpired: return None except OSError: return subprocess.CompletedProcess(command, 127, b"") process = run_worker(limit) if wide and (process is None or process.returncode == 124): # Wide documents can form one prohibitively expensive layout block. # The three generic views cover the same uploaded page without a # document-specific field template. process = run_worker(90, tiles_only=True) elapsed = round((time.monotonic() - started) * 1000) if process is None or process.returncode == 124: return {"status": "error", "error": "vision_timeout", "elapsed_ms": elapsed} if process.returncode or len(process.stdout) > 2_000_000: return {"status": "error", "error": "vision_failed", "elapsed_ms": elapsed} try: raw = json.loads(process.stdout) if raw.get("error"): code = "vision_models_missing" if raw["error"] == "models_missing" else "vision_failed" return {"status": "error", "error": code, "elapsed_ms": elapsed} lines = raw["lines"] regions = raw["regions"] if not isinstance(lines, list) or not isinstance(regions, list): raise ValueError("invalid collections") except (ValueError, TypeError, KeyError): return {"status": "error", "error": "vision_invalid", "elapsed_ms": elapsed} return {"status": "ok", "model": raw.get("model", "PaddleOCR-VL"), "elapsed_ms": elapsed, "lines": lines, "fields": assign_fields(lines, allow_adjacent=False), "regions": regions, "layout_note": raw.get("layout_note", ""), "retry_used": bool(raw.get("retry_used", False))} async def index(_: Request): return FileResponse(HERE / "index.html", media_type="text/html; charset=utf-8") async def script(_: Request): return FileResponse(HERE / "app.js", media_type="text/javascript; charset=utf-8") async def style(_: Request): return FileResponse(HERE / "style.css", media_type="text/css; charset=utf-8") async def health(_: Request): return JSONResponse({"ready": True, "scope": "localhost only", "vision_mode": os.environ.get("OCR_VISION_MODE", "local")}) async def compare(request: Request): if _slot.locked(): return JSONResponse({"error": "busy"}, status_code=429) if request.headers.get("content-type", "").split(";", 1)[0] not in ("image/jpeg", "image/png"): return JSONResponse({"error": "unsupported_media_type"}, status_code=415) detail = request.headers.get("x-ocr-detail", "standard") if detail not in ("standard", "fine"): return JSONResponse({"error": "invalid_detail"}, status_code=422) async with _slot: data = bytearray() try: async with asyncio.timeout(30): async for chunk in request.stream(): data.extend(chunk) if len(data) > MAX_UPLOAD: return JSONResponse({"error": "upload_too_large"}, status_code=413) except TimeoutError: return JSONResponse({"error": "upload_timeout"}, status_code=408) try: image, (width, height) = await asyncio.to_thread( _prepare, bytes(data), request.headers.get("x-ocr-geometry")) except InputError as exc: return JSONResponse({"error": str(exc)}, status_code=422) with tempfile.TemporaryDirectory(prefix="ocr-compare-") as path: directory = Path(path) classic, vision = await asyncio.gather( asyncio.to_thread(_classic, image, directory, fine_text=detail == "fine"), asyncio.to_thread(_vision, image), ) # Exact JPEG bytes sent to both pipelines are returned for evidence # highlighting. They never enter a report, log, or external OCR API. import base64 return JSONResponse({"image": "data:image/jpeg;base64," + base64.b64encode(image).decode(), "image_size": {"width": width, "height": height}, "field_codes": FIELD_CODES, "results": {"classic": classic, "vision": vision}}) app = Starlette(routes=[Route("/", index), Route("/app.js", script), Route("/style.css", style), Route("/api/health", health), Route("/api/compare", compare, methods=["POST"])])