zulassung-ocr-compare/ocr_compare/server.py

"""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"])])