klarbildbot/app/enhance.py

"""AI-Bildverarbeitung: Freistellen (rembg) + Upscaling (dnn_superres / Real-ESRGAN).

Alle Operationen laufen lokal auf CPU. Keine externen/US-APIs.
"""
from __future__ import annotations

import io
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import cv2
import numpy as np
from PIL import Image

from . import config

# ---------------------------------------------------------------------------
# Lazy-geladene, wiederverwendete Modelle (teuer beim ersten Mal).
# ---------------------------------------------------------------------------
_rembg_session = None
_sr_cache: dict[tuple[str, int], "cv2.dnn_superres.DnnSuperResImpl"] = {}
_ort_session = None


@dataclass
class Result:
    data: bytes
    fmt: str          # "png" oder "jpg"
    width: int
    height: int
    seconds: float
    note: str = ""


# ---------------------------------------------------------------------------
# Hilfsfunktionen
# ---------------------------------------------------------------------------
def _load_bgr(image_bytes: bytes) -> np.ndarray:
    """Bytes -> BGR uint8 (EXIF-Rotation beruecksichtigt)."""
    pil = Image.open(io.BytesIO(image_bytes))
    try:
        from PIL import ImageOps

        pil = ImageOps.exif_transpose(pil)
    except Exception:
        pass
    pil = pil.convert("RGB")
    rgb = np.array(pil)
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)


def _limit_input(img: np.ndarray, max_edge: int) -> np.ndarray:
    h, w = img.shape[:2]
    longest = max(h, w)
    if longest <= max_edge:
        return img
    scale = max_edge / longest
    new = (int(round(w * scale)), int(round(h * scale)))
    return cv2.resize(img, new, interpolation=cv2.INTER_AREA)


def _encode(img_bgr_or_bgra: np.ndarray, fmt: str) -> bytes:
    if fmt == "png":
        ok, buf = cv2.imencode(".png", img_bgr_or_bgra,
                               [cv2.IMWRITE_PNG_COMPRESSION, 6])
    else:
        ok, buf = cv2.imencode(".jpg", img_bgr_or_bgra,
                               [cv2.IMWRITE_JPEG_QUALITY, config.JPEG_QUALITY])
    if not ok:
        raise RuntimeError("Bild-Encoding fehlgeschlagen")
    return buf.tobytes()


def _post_process(img: np.ndarray) -> np.ndarray:
    """Leichte Entrauschung + Unsharp-Mask fuer den 'Klarbild'-Look."""
    out = img
    if config.POST_DENOISE:
        out = cv2.fastNlMeansDenoisingColored(out, None, 3, 3, 7, 21)
    if config.POST_SHARPEN:
        blur = cv2.GaussianBlur(out, (0, 0), 1.2)
        out = cv2.addWeighted(out, 1.5, blur, -0.5, 0)
    return out


# ---------------------------------------------------------------------------
# Upscaling
# ---------------------------------------------------------------------------
def _get_sr_model(model: str, scale: int):
    key = (model, scale)
    if key in _sr_cache:
        return _sr_cache[key]
    fname = f"{model.upper()}_x{scale}.pb"
    path = config.MODELS_DIR / fname
    if not path.exists():
        raise FileNotFoundError(f"SR-Modell fehlt: {path}")
    sr = cv2.dnn_superres.DnnSuperResImpl_create()
    sr.readModel(str(path))
    sr.setModel(model.lower(), scale)
    _sr_cache[key] = sr
    return sr


def _tiled_sr(img: np.ndarray, sr, scale: int, tile: int) -> np.ndarray:
    """Kachelweises Upscaling mit Ueberlappung gegen Kachel-Kanten."""
    h, w = img.shape[:2]
    if tile <= 0 or (h <= tile and w <= tile):
        return sr.upsample(img)

    pad = 16  # Ueberlappung
    out = np.zeros((h * scale, w * scale, 3), dtype=np.uint8)
    for y in range(0, h, tile):
        for x in range(0, w, tile):
            y0, x0 = max(0, y - pad), max(0, x - pad)
            y1, x1 = min(h, y + tile + pad), min(w, x + tile + pad)
            patch = img[y0:y1, x0:x1]
            up = sr.upsample(patch)
            # gueltigen (nicht-ueberlappenden) Bereich zurueckschneiden
            ty0, tx0 = (y - y0) * scale, (x - x0) * scale
            vy, vx = min(tile, h - y), min(tile, w - x)
            crop = up[ty0:ty0 + vy * scale, tx0:tx0 + vx * scale]
            out[y * scale:y * scale + vy * scale,
                x * scale:x * scale + vx * scale] = crop
    return out


def _upscale_opencv(img: np.ndarray, scale: int) -> tuple[np.ndarray, str]:
    """scale in {2,4}. Nutzt konfiguriertes Modell, faellt auf fsrcnn zurueck."""
    model = config.UPSCALE_MODEL
    try:
        sr = _get_sr_model(model, scale)
        used = model
    except FileNotFoundError:
        sr = _get_sr_model("fsrcnn", scale)
        used = "fsrcnn"
    out = _tiled_sr(img, sr, scale, config.TILE_SIZE)
    return out, f"dnn_superres/{used}_x{scale}"


def _upscale_realesrgan(img: np.ndarray) -> tuple[np.ndarray, str]:
    """Real-ESRGAN x4 via onnxruntime (optional, fuer starke Boxen/GPU)."""
    global _ort_session
    import onnxruntime as ort  # lokal importieren

    if _ort_session is None:
        _ort_session = ort.InferenceSession(
            str(config.REALESRGAN_ONNX),
            providers=["CPUExecutionProvider"],
        )
    sess = _ort_session
    iname = sess.get_inputs()[0].name

    def run(patch_bgr: np.ndarray) -> np.ndarray:
        rgb = cv2.cvtColor(patch_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        inp = np.transpose(rgb, (2, 0, 1))[None, ...]
        out = sess.run(None, {iname: inp})[0]
        out = np.clip(out[0], 0, 1)
        out = np.transpose(out, (1, 2, 0))
        out = (out * 255.0).round().astype(np.uint8)
        return cv2.cvtColor(out, cv2.COLOR_RGB2BGR)

    h, w = img.shape[:2]
    tile, pad, scale = config.TILE_SIZE, 16, 4
    out = np.zeros((h * scale, w * scale, 3), dtype=np.uint8)
    for y in range(0, h, tile):
        for x in range(0, w, tile):
            y0, x0 = max(0, y - pad), max(0, x - pad)
            y1, x1 = min(h, y + tile + pad), min(w, x + tile + pad)
            up = run(img[y0:y1, x0:x1])
            ty0, tx0 = (y - y0) * scale, (x - x0) * scale
            vy, vx = min(tile, h - y), min(tile, w - x)
            crop = up[ty0:ty0 + vy * scale, tx0:tx0 + vx * scale]
            out[y * scale:y * scale + vy * scale,
                x * scale:x * scale + vx * scale] = crop
    return out, "real-esrgan_x4(onnx)"


def upscale(image_bytes: bytes, scale: int) -> Result:
    t0 = time.time()
    img = _load_bgr(image_bytes)
    img = _limit_input(img, config.MAX_INPUT_EDGE)
    h, w = img.shape[:2]

    # Ausgabe-Megapixel deckeln
    out_mp = (h * scale) * (w * scale) / 1_000_000
    note = ""
    if out_mp > config.MAX_OUTPUT_MP and scale == 4:
        scale = 2
        note = "Auf ×2 begrenzt (Speicherschutz)."

    use_realesrgan = (
        config.ENABLE_REALESRGAN and config.REALESRGAN_ONNX.exists()
    )
    if use_realesrgan and scale == 4:
        up, backend = _upscale_realesrgan(img)
    elif use_realesrgan and scale == 2:
        up, backend = _upscale_realesrgan(img)
        up = cv2.resize(up, (w * 2, h * 2), interpolation=cv2.INTER_AREA)
        backend = "real-esrgan_x4->x2"
    else:
        up, backend = _upscale_opencv(img, scale)

    up = _post_process(up)
    data = _encode(up, "jpg")
    oh, ow = up.shape[:2]
    return Result(data, "jpg", ow, oh, time.time() - t0,
                  (f"{backend} · {note}").strip(" ·"))


# ---------------------------------------------------------------------------
# Freistellen (Hintergrund entfernen)
# ---------------------------------------------------------------------------
def _get_rembg():
    global _rembg_session
    if _rembg_session is None:
        from rembg import new_session

        _rembg_session = new_session(config.REMBG_MODEL)
    return _rembg_session


def cutout(image_bytes: bytes) -> Result:
    """Hintergrund per rembg entfernen -> PNG mit Transparenz."""
    from rembg import remove

    t0 = time.time()
    session = _get_rembg()
    out_png = remove(
        image_bytes,
        session=session,
        alpha_matting=True,
        alpha_matting_foreground_threshold=240,
        alpha_matting_background_threshold=15,
        alpha_matting_erode_size=8,
    )
    pil = Image.open(io.BytesIO(out_png)).convert("RGBA")
    w, h = pil.size
    return Result(out_png, "png", w, h, time.time() - t0,
                  f"rembg/{config.REMBG_MODEL}")


def cutout_then_upscale(image_bytes: bytes, scale: int) -> Result:
    """Erst freistellen, dann das freigestellte Motiv hochskalieren (RGBA)."""
    t0 = time.time()
    cut = cutout(image_bytes)
    rgba = Image.open(io.BytesIO(cut.data)).convert("RGBA")
    alpha = np.array(rgba.split()[-1])
    rgb = cv2.cvtColor(np.array(rgba.convert("RGB")), cv2.COLOR_RGB2BGR)

    up = upscale(_encode(rgb, "png"), scale)  # nutzt gesamte Pipeline
    up_bgr = cv2.imdecode(np.frombuffer(up.data, np.uint8), cv2.IMREAD_COLOR)

    # Alpha passend hochskalieren und wieder anlegen
    oh, ow = up_bgr.shape[:2]
    alpha_up = cv2.resize(alpha, (ow, oh), interpolation=cv2.INTER_LINEAR)
    bgra = cv2.cvtColor(up_bgr, cv2.COLOR_BGR2BGRA)
    bgra[:, :, 3] = alpha_up
    data = _encode(bgra, "png")
    return Result(data, "png", ow, oh, time.time() - t0,
                  f"{cut.note} + {up.note}")


def warmup() -> None:
    """Modelle beim Start vorladen, damit die erste Anfrage schnell ist."""
    try:
        _get_rembg()
    except Exception as exc:  # pragma: no cover
        print(f"[warmup] rembg nicht geladen: {exc}", flush=True)
    for scale in (2, 4):
        try:
            _get_sr_model(config.UPSCALE_MODEL, scale)
        except Exception as exc:  # pragma: no cover
            print(f"[warmup] SR {config.UPSCALE_MODEL} x{scale}: {exc}", flush=True)