familienarchiv/ocr-service/main.py

"""OCR microservice — FastAPI app with Surya and Kraken engine support."""

import asyncio
import glob
import io
import json
import logging
import os
import shutil
import tempfile
import zipfile
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from urllib.parse import urlparse

import httpx
import pypdfium2 as pdfium
from fastapi import FastAPI, Header, HTTPException, UploadFile
from fastapi.responses import StreamingResponse
from PIL import Image

from confidence import apply_confidence_markers, get_threshold
from engines import kraken as kraken_engine
from engines import surya as surya_engine
from models import OcrBlock, OcrRequest

TRAINING_TOKEN = os.environ.get("TRAINING_TOKEN", "")
KRAKEN_MODEL_PATH = os.environ.get("KRAKEN_MODEL_PATH", "/app/models/german_kurrent.mlmodel")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

_models_ready = False

ALLOWED_PDF_HOSTS = set(
    h.strip() for h in os.getenv("ALLOWED_PDF_HOSTS", "minio,localhost,127.0.0.1").split(",")
)


def _validate_url(url: str) -> None:
    """Validate that the PDF URL points to an allowed host (SSRF protection)."""
    parsed = urlparse(url)
    hostname = parsed.hostname or ""
    if hostname not in ALLOWED_PDF_HOSTS:
        raise HTTPException(status_code=400, detail=f"PDF host not allowed: {hostname}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load lightweight models at startup. Surya loads lazily on first request."""
    global _models_ready

    logger.info("Loading Kraken model at startup (Surya loads lazily on first OCR request)...")
    kraken_engine.load_models()
    _models_ready = True
    logger.info("Startup complete — ready to accept requests")

    yield

    logger.info("Shutting down OCR service")


app = FastAPI(title="Familienarchiv OCR Service", lifespan=lifespan)


@app.get("/health")
def health():
    """Health endpoint — returns 200 only after models are loaded."""
    if not _models_ready:
        raise HTTPException(status_code=503, detail="Models not loaded yet")
    return {"status": "ok", "surya": True, "kraken": kraken_engine.is_available()}


@app.post("/ocr", response_model=list[OcrBlock])
async def run_ocr(request: OcrRequest):
    """Run OCR on a PDF document.

    Downloads the PDF from the provided URL, converts pages to images,
    and runs the appropriate OCR engine based on scriptType.
    OCR engines run in a thread pool so the event loop stays free for /health.
    """
    if not _models_ready:
        raise HTTPException(status_code=503, detail="Models not loaded yet")

    images = await _download_and_convert_pdf(request.pdfUrl)

    script_type = request.scriptType.upper()

    if script_type == "HANDWRITING_KURRENT":
        if not kraken_engine.is_available():
            raise HTTPException(
                status_code=400,
                detail="Kraken model not available — cannot process Kurrent script",
            )
        blocks = await asyncio.to_thread(kraken_engine.extract_blocks, images, request.language)
    else:
        # TYPEWRITER, HANDWRITING_LATIN, UNKNOWN — all use Surya
        blocks = await asyncio.to_thread(surya_engine.extract_blocks, images, request.language)

    threshold = get_threshold(script_type)
    for block in blocks:
        if block.get("words"):
            block["text"] = apply_confidence_markers(block["words"], threshold)
        block.pop("words", None)

    return [OcrBlock(**b) for b in blocks]


@app.post("/ocr/stream")
async def run_ocr_stream(request: OcrRequest):
    """Run OCR on a PDF with NDJSON streaming — one JSON line per completed page."""
    if not _models_ready:
        raise HTTPException(status_code=503, detail="Models not loaded yet")

    images = await _download_and_convert_pdf(request.pdfUrl)
    script_type = request.scriptType.upper()
    threshold = get_threshold(script_type)

    use_kraken = script_type == "HANDWRITING_KURRENT"
    if use_kraken and not kraken_engine.is_available():
        raise HTTPException(
            status_code=400,
            detail="Kraken model not available — cannot process Kurrent script",
        )

    async def generate():
        total_pages = len(images)
        yield json.dumps({"type": "start", "totalPages": total_pages}) + "\n"

        total_blocks = 0
        skipped_pages = 0

        for page_idx, image in enumerate(images, start=1):
            try:
                engine = kraken_engine if use_kraken else surya_engine
                blocks = await asyncio.to_thread(
                    engine.extract_page_blocks, image, page_idx, request.language
                )

                for block in blocks:
                    if block.get("words"):
                        block["text"] = apply_confidence_markers(block["words"], threshold)
                    block.pop("words", None)

                total_blocks += len(blocks)
                yield json.dumps({
                    "type": "page",
                    "pageNumber": page_idx,
                    "blocks": blocks,
                }) + "\n"

            except Exception:
                logger.exception("OCR failed on page %d", page_idx)
                skipped_pages += 1
                yield json.dumps({
                    "type": "error",
                    "pageNumber": page_idx,
                    "message": f"OCR processing failed on page {page_idx}",
                }) + "\n"

            finally:
                del image

        yield json.dumps({
            "type": "done",
            "totalBlocks": total_blocks,
            "skippedPages": skipped_pages,
        }) + "\n"

    return StreamingResponse(
        generate(),
        media_type="application/x-ndjson",
        headers={
            "X-Accel-Buffering": "no",
            "Cache-Control": "no-cache",
        },
    )


def _check_training_token(x_training_token: str | None) -> None:
    """Validate training token if TRAINING_TOKEN env var is set."""
    if TRAINING_TOKEN and x_training_token != TRAINING_TOKEN:
        raise HTTPException(status_code=403, detail="Invalid or missing X-Training-Token")


def _validate_zip_entry(name: str, extract_dir: str) -> None:
    """Reject ZIP Slip attacks: path traversal and absolute paths."""
    if os.path.isabs(name) or name.startswith(".."):
        raise HTTPException(status_code=400, detail=f"Unsafe ZIP entry: {name}")
    resolved = os.path.realpath(os.path.join(extract_dir, name))
    if not resolved.startswith(os.path.realpath(extract_dir)):
        raise HTTPException(status_code=400, detail=f"ZIP Slip detected: {name}")


def _rotate_backups(model_path: str, keep: int = 3) -> None:
    """Keep only the last `keep` timestamped backups of the model."""
    pattern = model_path + ".*.bak"
    backups = sorted(glob.glob(pattern))
    for old in backups[:-keep]:
        try:
            os.remove(old)
        except OSError:
            logger.warning("Could not remove old backup: %s", old)


@app.post("/train")
async def train_model(
    file: UploadFile,
    x_training_token: str | None = Header(default=None),
):
    """Fine-tune the Kurrent recognition model with uploaded training data.

    Accepts a ZIP archive containing .png/.gt.txt training pairs exported
    by the Java backend. Training mutates in-process model state — not safe
    if the service is replicated.
    """
    _check_training_token(x_training_token)

    if not _models_ready:
        raise HTTPException(status_code=503, detail="Models not loaded yet")

    zip_bytes = await file.read()

    training_run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    log = logging.LoggerAdapter(logger, {"training_run_id": training_run_id})
    log.info("Starting training run %s", training_run_id)

    def _run_training() -> dict:
        with tempfile.TemporaryDirectory() as tmp_dir:
            # Extract ZIP with safety checks
            with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
                for entry in zf.namelist():
                    _validate_zip_entry(entry, tmp_dir)
                zf.extractall(tmp_dir)

            log.info("Extracted %d ZIP entries to %s", len(os.listdir(tmp_dir)), tmp_dir)

            # Run ketos train (transfer learning from existing model)
            from kraken import ketos
            ground_truth = glob.glob(os.path.join(tmp_dir, "*.gt.txt"))
            if not ground_truth:
                raise HTTPException(status_code=422, detail="No ground-truth files found in ZIP")

            log.info("Training on %d ground-truth pairs", len(ground_truth))
            output_model_path = os.path.join(tmp_dir, "fine_tuned.mlmodel")

            result = ketos.train(
                ground_truth=ground_truth,
                load=KRAKEN_MODEL_PATH,
                output=output_model_path,
                format_type="path",
            )

            epochs = getattr(result, "epochs", None) or 0
            loss = getattr(result, "best_loss", None)
            accuracy = getattr(result, "best_accuracy", None)

            log.info("Training complete — epochs=%s loss=%s accuracy=%s", epochs, loss, accuracy)

            # Backup existing model and replace
            if os.path.exists(KRAKEN_MODEL_PATH):
                timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
                backup_path = f"{KRAKEN_MODEL_PATH}.{timestamp}.bak"
                shutil.copy2(KRAKEN_MODEL_PATH, backup_path)
                log.info("Backed up model to %s", backup_path)
                _rotate_backups(KRAKEN_MODEL_PATH, keep=3)

            shutil.move(output_model_path, KRAKEN_MODEL_PATH)
            log.info("Replaced model at %s", KRAKEN_MODEL_PATH)

            # Reload model in-process
            kraken_engine.load_models()
            log.info("Reloaded Kraken model in-process")

            return {"loss": loss, "accuracy": accuracy, "epochs": epochs}

    result = await asyncio.to_thread(_run_training)
    return result


@app.post("/segtrain")
async def segtrain_model(
    file: UploadFile,
    x_training_token: str | None = Header(default=None),
):
    """Fine-tune the blla segmentation model with uploaded PAGE XML training data.

    Accepts a ZIP archive containing .png/.xml (PAGE XML) training pairs exported
    by the Java backend. Training mutates in-process model state — not safe
    if the service is replicated.
    """
    _check_training_token(x_training_token)

    if not _models_ready:
        raise HTTPException(status_code=503, detail="Models not loaded yet")

    zip_bytes = await file.read()

    training_run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    log = logging.LoggerAdapter(logger, {"training_run_id": training_run_id})
    log.info("Starting segmentation training run %s", training_run_id)

    blla_model_path = os.environ.get("BLLA_MODEL_PATH", "/app/models/blla.mlmodel")

    def _run_segtrain() -> dict:
        with tempfile.TemporaryDirectory() as tmp_dir:
            with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
                for entry in zf.namelist():
                    _validate_zip_entry(entry, tmp_dir)
                zf.extractall(tmp_dir)

            log.info("Extracted %d ZIP entries for segmentation training", len(os.listdir(tmp_dir)))

            xml_files = glob.glob(os.path.join(tmp_dir, "*.xml"))
            if not xml_files:
                raise HTTPException(status_code=422, detail="No PAGE XML files found in ZIP")

            log.info("Training on %d PAGE XML files", len(xml_files))
            output_model_path = os.path.join(tmp_dir, "fine_tuned_blla.mlmodel")

            from kraken import ketos
            result = ketos.segtrain(
                ground_truth=xml_files,
                load=blla_model_path if os.path.exists(blla_model_path) else None,
                output=output_model_path,
                format_type="path",
            )

            epochs = getattr(result, "epochs", None) or 0
            loss = getattr(result, "best_loss", None)
            accuracy = getattr(result, "best_accuracy", None)

            log.info("Segmentation training complete — epochs=%s loss=%s", epochs, loss)

            if os.path.exists(blla_model_path):
                timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
                backup_path = f"{blla_model_path}.{timestamp}.bak"
                shutil.copy2(blla_model_path, backup_path)
                _rotate_backups(blla_model_path, keep=3)

            shutil.move(output_model_path, blla_model_path)
            log.info("Replaced blla model at %s", blla_model_path)

            return {"loss": loss, "accuracy": accuracy, "epochs": epochs}

    result = await asyncio.to_thread(_run_segtrain)
    return result


async def _download_and_convert_pdf(url: str) -> list[Image.Image]:
    """Download a PDF from a presigned URL and convert each page to a PIL Image."""
    _validate_url(url)
    async with httpx.AsyncClient(
        timeout=httpx.Timeout(300.0), follow_redirects=False
    ) as client:
        response = await client.get(url)
        response.raise_for_status()

    pdf = pdfium.PdfDocument(io.BytesIO(response.content))
    images = []

    for page_idx in range(len(pdf)):
        page = pdf[page_idx]
        # Render at 200 DPI — balances OCR quality vs memory usage
        # (Surya 0.17 models use ~5GB idle; 300 DPI causes OOM on multi-page docs)
        bitmap = page.render(scale=200 / 72)
        pil_image = bitmap.to_pil()
        images.append(pil_image)

    return images