PDF-Qualität mit KI verbessern: VLMs und OCR auf Kubernetes

# preprocessing.py
import fitz  # PyMuPDF
from PIL import Image, ImageFilter, ImageEnhance
import numpy as np

def extract_page_as_image(pdf_path: str, page_num: int, zoom: float = 2.0) -> Image.Image:
    """Extrahiert eine PDF-Seite als hochaufgeloestes Bild."""
    doc = fitz.open(pdf_path)
    page = doc.load_page(page_num)
    mat = fitz.Matrix(zoom, zoom)
    pix = page.get_pixmap(matrix=mat)
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    doc.close()
    return img

def enhance_for_ocr(image: Image.Image) -> Image.Image:
    """Optimiert ein Bild fuer nachfolgende OCR/VLM-Verarbeitung."""
    # Kontrast erhoehen
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(1.5)

    # Schaerfen
    image = image.filter(ImageFilter.SHARPEN)

    # In Graustufen konvertieren und zurueck (reduziert Rauschen)
    gray = image.convert("L")
    # Adaptive Binarisierung per Schwellwert
    threshold = np.mean(np.array(gray))
    binary = gray.point(lambda x: 255 if x > threshold else 0)

    return binary.convert("RGB")

# vlm_processor.py
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from transformers import LayoutLMv3Processor, LayoutLMv3ForSequenceClassification
import torch
from PIL import Image
from typing import List, Dict

class DocumentProcessor:
    def __init__(self, device: str = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        # TrOCR fuer Texterkennung
        self.trocr_processor = TrOCRProcessor.from_pretrained(
            "microsoft/trocr-base-handwritten"
        )
        self.trocr_model = VisionEncoderDecoderModel.from_pretrained(
            "microsoft/trocr-base-handwritten"
        ).to(self.device)

        # LayoutLMv3 fuer Dokumentenverstaendnis
        self.layout_processor = LayoutLMv3Processor.from_pretrained(
            "microsoft/layoutlmv3-base"
        )
        self.layout_model = LayoutLMv3ForSequenceClassification.from_pretrained(
            "microsoft/layoutlmv3-base"
        ).to(self.device)

    def extract_text(self, image: Image.Image) -> str:
        """Extrahiert Text aus einem Bild per TrOCR."""
        pixel_values = self.trocr_processor(
            image, return_tensors="pt"
        ).pixel_values.to(self.device)

        with torch.no_grad():
            generated_ids = self.trocr_model.generate(pixel_values, max_new_tokens=512)

        return self.trocr_processor.batch_decode(
            generated_ids, skip_special_tokens=True
        )[0]

    def classify_regions(self, image: Image.Image) -> List[Dict]:
        """Klassifiziert Regionen eines Dokuments (Header, Table, Text, etc.)."""
        encoding = self.layout_processor(
            image, return_tensors="pt", truncation=True
        )
        encoding = {k: v.to(self.device) for k, v in encoding.items()}

        with torch.no_grad():
            outputs = self.layout_model(**encoding)
            predictions = torch.softmax(outputs.logits, dim=1)

        labels = ["header", "footer", "title", "text", "table",
                  "figure", "list", "form", "signature"]

        return [
            {"label": labels[i], "confidence": predictions[0][i].item()}
            for i in range(len(labels))
            if predictions[0][i].item() > 0.1
        ]

# pdf-processor-deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: pdf-processor
  namespace: document-ai
spec:
  replicas: 2
  selector:
    matchLabels:
      app: pdf-processor
  template:
    metadata:
      labels:
        app: pdf-processor
    spec:
      containers:
        - name: processor
          image: registry.example.com/pdf-processor:v1.2.0
          ports:
            - containerPort: 8000
          resources:
            requests:
              cpu: "2"
              memory: "8Gi"
              nvidia.com/gpu: "1"
            limits:
              cpu: "4"
              memory: "16Gi"
              nvidia.com/gpu: "1"
          env:
            - name: MODEL_CACHE_DIR
              value: "/models"
            - name: WORKERS
              value: "2"
          volumeMounts:
            - name: model-cache
              mountPath: /models
            - name: upload-volume
              mountPath: /data/uploads
          readinessProbe:
            httpGet:
              path: /health
              port: 8000
            initialDelaySeconds: 30
            periodSeconds: 10
          livenessProbe:
            httpGet:
              path: /health
              port: 8000
            initialDelaySeconds: 60
            periodSeconds: 30
      volumes:
        - name: model-cache
          persistentVolumeClaim:
            claimName: model-cache-pvc
        - name: upload-volume
          persistentVolumeClaim:
            claimName: upload-pvc
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
      nodeSelector:
        accelerator: nvidia-a100
---
apiVersion: v1
kind: Service
metadata:
  name: pdf-processor
  namespace: document-ai
spec:
  selector:
    app: pdf-processor
  ports:
    - port: 80
      targetPort: 8000
  type: ClusterIP
---
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: pdf-processor-hpa
  namespace: document-ai
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: pdf-processor
  minReplicas: 1
  maxReplicas: 6
  metrics:
    - type: Pods
      pods:
        metric:
          name: queue_depth
        target:
          type: AverageValue
          averageValue: "5"

# api.py
from fastapi import FastAPI, UploadFile, File, BackgroundTasks
from fastapi.responses import JSONResponse
import uuid
import os
import redis
import json

app = FastAPI(title="PDF Quality Enhancement API")
cache = redis.Redis(host=os.getenv("REDIS_HOST", "redis"), port=6379, db=0)

processor = DocumentProcessor()

@app.post("/process")
async def process_pdf(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    """Nimmt ein PDF entgegen und startet asynchrone Verarbeitung."""
    task_id = str(uuid.uuid4())
    upload_path = f"/data/uploads/{task_id}.pdf"

    with open(upload_path, "wb") as f:
        content = await file.read()
        f.write(content)

    cache.set(f"task:{task_id}", json.dumps({"status": "queued"}))
    background_tasks.add_task(run_pipeline, upload_path, task_id)

    return {"task_id": task_id, "status": "queued"}

@app.get("/status/{task_id}")
async def get_status(task_id: str):
    """Gibt den Verarbeitungsstatus zurueck."""
    result = cache.get(f"task:{task_id}")
    if not result:
        return JSONResponse(status_code=404, content={"error": "Task nicht gefunden"})
    return json.loads(result)

@app.get("/health")
async def health():
    return {"status": "ok", "gpu_available": torch.cuda.is_available()}

async def run_pipeline(pdf_path: str, task_id: str):
    """Verarbeitet ein PDF durch die gesamte Pipeline."""
    try:
        cache.set(f"task:{task_id}", json.dumps({"status": "processing"}))

        from preprocessing import extract_page_as_image, enhance_for_ocr

        doc = fitz.open(pdf_path)
        results = []

        for page_num in range(len(doc)):
            image = extract_page_as_image(pdf_path, page_num)
            enhanced = enhance_for_ocr(image)
            text = processor.extract_text(enhanced)
            regions = processor.classify_regions(enhanced)

            results.append({
                "page": page_num + 1,
                "text": text,
                "regions": regions
            })

        doc.close()
        output = {"status": "completed", "pages": len(results), "results": results}
        cache.set(f"task:{task_id}", json.dumps(output))

    except Exception as e:
        cache.set(f"task:{task_id}", json.dumps({"status": "failed", "error": str(e)}))