Klarbildbot v1.0 — AI-Bildaufbereitung (Freistellen + Upscaling) als Telegram-Bot
- Telegram-Bot (Polling) mit Inline-Menue: Klarbild x2/x4, Freistellen, kombiniert - Upscaling via OpenCV dnn_superres (FSRCNN default, EDSR optional), gekachelt - Freistellen via rembg (isnet-general-use, Alpha-Matting) - Picdrop-Batch via SFTP (/picdrop) - Health-Server (aiohttp) fuer Coolify, Zugriffsschutz via ALLOWED_USER_IDS - Dockerfile backt Modelle (SR von raw.githubusercontent, rembg von HF-Mirror) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XNQ8ghPfzAfsyVYd6HgFb6
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
"""Klarbildbot – AI-Bildaufbereitung (Freistellen + Upscaling) als Telegram-Bot."""
|
||||
__version__ = "1.0.0"
|
||||
+344
@@ -0,0 +1,344 @@
|
||||
"""Klarbildbot – Telegram-Bot fuer AI-Bildaufbereitung (Freistellen + Upscaling)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import posixpath
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
from telegram import (
|
||||
InlineKeyboardButton,
|
||||
InlineKeyboardMarkup,
|
||||
InputFile,
|
||||
Update,
|
||||
)
|
||||
from telegram.constants import ChatAction
|
||||
from telegram.ext import (
|
||||
Application,
|
||||
CallbackQueryHandler,
|
||||
CommandHandler,
|
||||
ContextTypes,
|
||||
MessageHandler,
|
||||
filters,
|
||||
)
|
||||
|
||||
from . import config, enhance, picdrop
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s %(levelname)s %(name)s | %(message)s",
|
||||
level=logging.INFO,
|
||||
)
|
||||
log = logging.getLogger("klarbildbot")
|
||||
|
||||
# Schwere CPU-Jobs strikt serialisieren (Speicherschutz auf kleinen Boxen).
|
||||
EXECUTOR = ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
# Kurzzeit-Speicher: pending-Bild pro User (file_id + Herkunft).
|
||||
PENDING: dict[int, dict] = {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Zugriffsschutz
|
||||
# ---------------------------------------------------------------------------
|
||||
def _allowed(update: Update) -> bool:
|
||||
if not config.ALLOWED_USER_IDS:
|
||||
return True
|
||||
user = update.effective_user
|
||||
return bool(user and user.id in config.ALLOWED_USER_IDS)
|
||||
|
||||
|
||||
async def _deny(update: Update) -> None:
|
||||
uid = update.effective_user.id if update.effective_user else "?"
|
||||
await update.effective_message.reply_text(
|
||||
f"⛔️ Kein Zugriff. Deine Telegram-ID: {uid}\n"
|
||||
"Bitte vom Betreiber in ALLOWED_USER_IDS freischalten lassen."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Commands
|
||||
# ---------------------------------------------------------------------------
|
||||
async def cmd_start(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
uid = update.effective_user.id if update.effective_user else "?"
|
||||
if not _allowed(update):
|
||||
return await _deny(update)
|
||||
await update.effective_message.reply_text(
|
||||
"👋 *Klarbildbot* – AI-Bildaufbereitung\n\n"
|
||||
"Schick mir ein Bild, dann kannst du wählen:\n"
|
||||
"• 🔍 *Klarbild ×2 / ×4* – AI-Upscaling + Schärfen\n"
|
||||
"• ✂️ *Freistellen* – Hintergrund per AI entfernen (PNG)\n"
|
||||
"• ✨ *Freistellen + Klarbild* – beides\n\n"
|
||||
"💡 Für beste Qualität das Bild als *Datei* senden "
|
||||
"(Büroklammer → Datei), nicht als komprimiertes Foto.\n\n"
|
||||
f"Deine Telegram-ID: `{uid}`",
|
||||
parse_mode="Markdown",
|
||||
)
|
||||
|
||||
|
||||
async def cmd_help(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
if not _allowed(update):
|
||||
return await _deny(update)
|
||||
txt = (
|
||||
"*So geht's:*\n"
|
||||
"1. Bild senden (am besten als Datei)\n"
|
||||
"2. Aktion wählen\n"
|
||||
"3. Ergebnis kommt als Datei zurück\n\n"
|
||||
"*Befehle:*\n"
|
||||
"/start – Übersicht\n"
|
||||
"/help – diese Hilfe\n"
|
||||
)
|
||||
if config.PICDROP_ENABLED:
|
||||
txt += "/picdrop – Picdrop-Galerie stapelweise verarbeiten\n"
|
||||
await update.effective_message.reply_text(txt, parse_mode="Markdown")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bild empfangen -> Aktionsmenue
|
||||
# ---------------------------------------------------------------------------
|
||||
def _menu() -> InlineKeyboardMarkup:
|
||||
return InlineKeyboardMarkup([
|
||||
[InlineKeyboardButton("🔍 Klarbild ×2", callback_data="up:2"),
|
||||
InlineKeyboardButton("🔍 Klarbild ×4", callback_data="up:4")],
|
||||
[InlineKeyboardButton("✂️ Freistellen", callback_data="cut:0")],
|
||||
[InlineKeyboardButton("✨ Freistellen + Klarbild ×2", callback_data="cutup:2")],
|
||||
])
|
||||
|
||||
|
||||
async def on_image(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
if not _allowed(update):
|
||||
return await _deny(update)
|
||||
msg = update.effective_message
|
||||
|
||||
if msg.photo:
|
||||
file_id = msg.photo[-1].file_id
|
||||
source = "photo"
|
||||
elif msg.document and (msg.document.mime_type or "").startswith("image/"):
|
||||
file_id = msg.document.file_id
|
||||
source = "document"
|
||||
else:
|
||||
return await msg.reply_text("Bitte ein Bild senden.")
|
||||
|
||||
PENDING[update.effective_user.id] = {"file_id": file_id, "source": source}
|
||||
hint = "" if source == "document" else "\n_(Tipp: als Datei senden = volle Auflösung)_"
|
||||
await msg.reply_text(
|
||||
"Was soll ich damit machen?" + hint,
|
||||
reply_markup=_menu(),
|
||||
parse_mode="Markdown",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Button gedrueckt -> verarbeiten
|
||||
# ---------------------------------------------------------------------------
|
||||
async def on_action(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
query = update.callback_query
|
||||
await query.answer()
|
||||
if not _allowed(update):
|
||||
return await query.edit_message_text("⛔️ Kein Zugriff.")
|
||||
|
||||
uid = update.effective_user.id
|
||||
pend = PENDING.get(uid)
|
||||
if not pend:
|
||||
return await query.edit_message_text(
|
||||
"Kein Bild gemerkt – bitte schick es nochmal.")
|
||||
|
||||
action, _, arg = query.data.partition(":")
|
||||
scale = int(arg) if arg.isdigit() else 0
|
||||
|
||||
labels = {"up": f"Klarbild ×{scale}", "cut": "Freistellen",
|
||||
"cutup": f"Freistellen + Klarbild ×{scale}"}
|
||||
await query.edit_message_text(f"⏳ {labels.get(action, action)} läuft … "
|
||||
"(kann bei großen Bildern etwas dauern)")
|
||||
|
||||
# Original holen (volle Aufloesung)
|
||||
try:
|
||||
tg_file = await ctx.bot.get_file(pend["file_id"])
|
||||
image_bytes = bytes(await tg_file.download_as_bytearray())
|
||||
except Exception as exc:
|
||||
log.exception("Download fehlgeschlagen")
|
||||
return await query.edit_message_text(
|
||||
f"❌ Konnte das Bild nicht laden: {exc}\n"
|
||||
"(Telegram-Bots können Dateien bis 20 MB laden.)")
|
||||
|
||||
await ctx.bot.send_chat_action(update.effective_chat.id, ChatAction.UPLOAD_DOCUMENT)
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def work() -> enhance.Result:
|
||||
if action == "up":
|
||||
return enhance.upscale(image_bytes, scale)
|
||||
if action == "cut":
|
||||
return enhance.cutout(image_bytes)
|
||||
if action == "cutup":
|
||||
return enhance.cutout_then_upscale(image_bytes, scale)
|
||||
raise ValueError(action)
|
||||
|
||||
try:
|
||||
result = await loop.run_in_executor(EXECUTOR, work)
|
||||
except Exception as exc:
|
||||
log.exception("Verarbeitung fehlgeschlagen")
|
||||
return await query.edit_message_text(f"❌ Fehler bei der Verarbeitung: {exc}")
|
||||
|
||||
fname = f"klarbild.{result.fmt}"
|
||||
caption = (f"✅ {labels.get(action, action)}\n"
|
||||
f"{result.width}×{result.height}px · {result.seconds:.1f}s"
|
||||
+ (f"\n{result.note}" if result.note else ""))
|
||||
try:
|
||||
await ctx.bot.send_document(
|
||||
chat_id=update.effective_chat.id,
|
||||
document=InputFile(result.data, filename=fname),
|
||||
caption=caption,
|
||||
)
|
||||
await query.edit_message_text(f"✅ Fertig: {labels.get(action, action)}")
|
||||
except Exception as exc:
|
||||
log.exception("Senden fehlgeschlagen")
|
||||
await query.edit_message_text(f"❌ Konnte Ergebnis nicht senden: {exc}")
|
||||
finally:
|
||||
PENDING.pop(uid, None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Picdrop-Batch
|
||||
# ---------------------------------------------------------------------------
|
||||
async def cmd_picdrop(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
if not _allowed(update):
|
||||
return await _deny(update)
|
||||
if not config.PICDROP_ENABLED:
|
||||
return await update.effective_message.reply_text(
|
||||
"Picdrop ist nicht konfiguriert (PICDROP_USER / PICDROP_PASSWORD).")
|
||||
|
||||
args = ctx.args or []
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
if not args:
|
||||
try:
|
||||
entries = await loop.run_in_executor(EXECUTOR, picdrop.list_dir, "/")
|
||||
except Exception as exc:
|
||||
return await update.effective_message.reply_text(f"❌ SFTP-Fehler: {exc}")
|
||||
folders = [e["name"] for e in entries if e["is_dir"]][:40]
|
||||
listing = "\n".join(f"• {f}" for f in folders) or "(keine Ordner)"
|
||||
return await update.effective_message.reply_text(
|
||||
"*Picdrop – Ordner im Wurzelverzeichnis:*\n" + listing +
|
||||
"\n\nNutzung: `/picdrop <Ordnername-Teil> up2|up4|cut`",
|
||||
parse_mode="Markdown")
|
||||
|
||||
needle = args[0]
|
||||
op = (args[1] if len(args) > 1 else "up2").lower()
|
||||
|
||||
await update.effective_message.reply_text(f"🔎 Suche Galerie „{needle}“ …")
|
||||
try:
|
||||
matches = await loop.run_in_executor(EXECUTOR, picdrop.find_gallery, needle)
|
||||
except Exception as exc:
|
||||
return await update.effective_message.reply_text(f"❌ SFTP-Fehler: {exc}")
|
||||
if not matches:
|
||||
return await update.effective_message.reply_text("Keine passende Galerie gefunden.")
|
||||
|
||||
gallery = matches[0]
|
||||
images = await loop.run_in_executor(EXECUTOR, picdrop.list_images, gallery)
|
||||
if not images:
|
||||
return await update.effective_message.reply_text(
|
||||
f"Galerie `{gallery}` enthält keine Bilder.", parse_mode="Markdown")
|
||||
|
||||
await update.effective_message.reply_text(
|
||||
f"📁 `{gallery}` · {len(images)} Bilder · Aktion `{op}`\n"
|
||||
"Verarbeite und lade in Unterordner `klarbild/` hoch …",
|
||||
parse_mode="Markdown")
|
||||
|
||||
def process_one(remote_path: str) -> str:
|
||||
data = picdrop.download(remote_path)
|
||||
if op == "cut":
|
||||
res = enhance.cutout(data)
|
||||
elif op == "up4":
|
||||
res = enhance.upscale(data, 4)
|
||||
else:
|
||||
res = enhance.upscale(data, 2)
|
||||
base = posixpath.splitext(posixpath.basename(remote_path))[0]
|
||||
out_name = f"{base}_klar.{res.fmt}"
|
||||
out_path = posixpath.join(gallery, "klarbild", out_name)
|
||||
picdrop.upload(out_path, res.data)
|
||||
return out_name
|
||||
|
||||
done = 0
|
||||
for img in images:
|
||||
try:
|
||||
name = await loop.run_in_executor(EXECUTOR, process_one, img["path"])
|
||||
done += 1
|
||||
if done % 5 == 0 or done == len(images):
|
||||
await update.effective_message.reply_text(
|
||||
f"… {done}/{len(images)} fertig (zuletzt: {name})")
|
||||
except Exception as exc:
|
||||
log.exception("Picdrop-Bild fehlgeschlagen")
|
||||
await update.effective_message.reply_text(
|
||||
f"⚠️ {img['name']}: {exc}")
|
||||
|
||||
await update.effective_message.reply_text(
|
||||
f"✅ Fertig: {done}/{len(images)} Bilder → `{gallery}/klarbild/`",
|
||||
parse_mode="Markdown")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Health-Server (fuer Coolify) + App-Setup
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _start_health(app: Application) -> None:
|
||||
from aiohttp import web
|
||||
|
||||
async def health(_req):
|
||||
return web.json_response({"status": "ok", "service": "klarbildbot"})
|
||||
|
||||
server = web.Application()
|
||||
server.router.add_get("/", health)
|
||||
server.router.add_get("/health", health)
|
||||
runner = web.AppRunner(server)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "0.0.0.0", config.HEALTH_PORT)
|
||||
await site.start()
|
||||
app.bot_data["health_runner"] = runner
|
||||
log.info("Health-Server auf :%s", config.HEALTH_PORT)
|
||||
|
||||
|
||||
async def _post_init(app: Application) -> None:
|
||||
await _start_health(app)
|
||||
# Bot-Kommandos in Telegram-UI registrieren
|
||||
from telegram import BotCommand
|
||||
|
||||
cmds = [BotCommand("start", "Übersicht"), BotCommand("help", "Hilfe")]
|
||||
if config.PICDROP_ENABLED:
|
||||
cmds.append(BotCommand("picdrop", "Picdrop-Galerie verarbeiten"))
|
||||
await app.bot.set_my_commands(cmds)
|
||||
# Modelle vorwaermen (im Thread, blockiert den Loop nicht)
|
||||
asyncio.get_running_loop().run_in_executor(EXECUTOR, enhance.warmup)
|
||||
log.info("Klarbildbot bereit.")
|
||||
|
||||
|
||||
async def _post_shutdown(app: Application) -> None:
|
||||
runner = app.bot_data.get("health_runner")
|
||||
if runner:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def build_app() -> Application:
|
||||
config.validate()
|
||||
app = (
|
||||
Application.builder()
|
||||
.token(config.BOT_TOKEN)
|
||||
.post_init(_post_init)
|
||||
.post_shutdown(_post_shutdown)
|
||||
.build()
|
||||
)
|
||||
app.add_handler(CommandHandler("start", cmd_start))
|
||||
app.add_handler(CommandHandler("help", cmd_help))
|
||||
app.add_handler(CommandHandler("picdrop", cmd_picdrop))
|
||||
app.add_handler(MessageHandler(
|
||||
filters.PHOTO | filters.Document.IMAGE, on_image))
|
||||
app.add_handler(CallbackQueryHandler(on_action))
|
||||
return app
|
||||
|
||||
|
||||
def main() -> None:
|
||||
app = build_app()
|
||||
log.info("Starte Polling …")
|
||||
app.run_polling(allowed_updates=Update.ALL_TYPES, drop_pending_updates=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Zentrale Konfiguration aus Umgebungsvariablen."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _get_bool(name: str, default: bool = False) -> bool:
|
||||
val = os.getenv(name)
|
||||
if val is None:
|
||||
return default
|
||||
return val.strip().lower() in {"1", "true", "yes", "on", "ja"}
|
||||
|
||||
|
||||
def _get_int(name: str, default: int) -> int:
|
||||
try:
|
||||
return int(os.getenv(name, str(default)))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
# --- Telegram ---------------------------------------------------------------
|
||||
BOT_TOKEN: str = os.getenv("BOT_TOKEN", "").strip()
|
||||
|
||||
# Kommagetrennte Liste erlaubter Telegram-User-IDs. Leer = jeder darf (nur
|
||||
# fuer Tests empfohlen). /start zeigt jedem seine ID, damit man hier eintraegt.
|
||||
_allowed = os.getenv("ALLOWED_USER_IDS", "").replace(";", ",")
|
||||
ALLOWED_USER_IDS: set[int] = {
|
||||
int(x) for x in (p.strip() for p in _allowed.split(",")) if x.strip().isdigit()
|
||||
}
|
||||
|
||||
# --- Bild-Pipeline ----------------------------------------------------------
|
||||
MODELS_DIR: Path = Path(os.getenv("MODELS_DIR", "/app/models"))
|
||||
WORK_DIR: Path = Path(os.getenv("WORK_DIR", "/tmp/klarbild"))
|
||||
|
||||
# rembg-Modell fuers Freistellen: isnet-general-use (gut), u2net, u2netp (leicht),
|
||||
# birefnet-general (beste Qualitaet, schwer). Wird beim Build vorgeladen.
|
||||
REMBG_MODEL: str = os.getenv("REMBG_MODEL", "isnet-general-use")
|
||||
|
||||
# OpenCV dnn_superres: "fsrcnn" (schnell, Default) oder "edsr" (mehr Detail,
|
||||
# aber auf CPU sehr langsam ~40s/Kachel -> nur fuer starke Boxen empfohlen).
|
||||
UPSCALE_MODEL: str = os.getenv("UPSCALE_MODEL", "fsrcnn").lower()
|
||||
|
||||
# Optionales Real-ESRGAN-ONNX-Backend (fuer GPU/starke Boxen). Wenn aktiv und
|
||||
# Modelldatei vorhanden, wird es statt OpenCV genutzt.
|
||||
ENABLE_REALESRGAN: bool = _get_bool("ENABLE_REALESRGAN", False)
|
||||
REALESRGAN_ONNX: Path = Path(
|
||||
os.getenv("REALESRGAN_ONNX", str(MODELS_DIR / "realesrgan_x4.onnx"))
|
||||
)
|
||||
|
||||
# Speicherschutz auf kleinen Boxen (CX33: 8GB, keine GPU):
|
||||
# Laengste Kante des Eingangsbildes vor dem Upscalen begrenzen.
|
||||
MAX_INPUT_EDGE: int = _get_int("MAX_INPUT_EDGE", 1600)
|
||||
# Kachelgroesse fuer das Upscalen (kleiner = weniger RAM, mehr Overhead).
|
||||
TILE_SIZE: int = _get_int("TILE_SIZE", 256)
|
||||
# Deckel fuer die Ausgabe-Megapixel (Sicherheitsnetz gegen OOM).
|
||||
MAX_OUTPUT_MP: float = float(os.getenv("MAX_OUTPUT_MP", "80"))
|
||||
|
||||
# Nach dem Upscalen leichte Entrauschung + Schaerfung ("Klarbild"-Finish).
|
||||
POST_DENOISE: bool = _get_bool("POST_DENOISE", True)
|
||||
POST_SHARPEN: bool = _get_bool("POST_SHARPEN", True)
|
||||
|
||||
# JPEG-Qualitaet fuer zurueckgesendete Fotos (Freistellen liefert immer PNG).
|
||||
JPEG_QUALITY: int = _get_int("JPEG_QUALITY", 95)
|
||||
|
||||
# --- Health-Server ----------------------------------------------------------
|
||||
HEALTH_PORT: int = _get_int("PORT", 8080)
|
||||
|
||||
# --- Picdrop (SFTP) ---------------------------------------------------------
|
||||
PICDROP_HOST: str = os.getenv("PICDROP_HOST", "ftps.picdrop.com")
|
||||
PICDROP_PORT: int = _get_int("PICDROP_PORT", 22)
|
||||
PICDROP_USER: str = os.getenv("PICDROP_USER", "")
|
||||
PICDROP_PASSWORD: str = os.getenv("PICDROP_PASSWORD", "")
|
||||
PICDROP_ENABLED: bool = bool(PICDROP_USER and PICDROP_PASSWORD)
|
||||
|
||||
|
||||
def validate() -> None:
|
||||
if not BOT_TOKEN:
|
||||
raise SystemExit("FEHLER: Umgebungsvariable BOT_TOKEN ist nicht gesetzt.")
|
||||
WORK_DIR.mkdir(parents=True, exist_ok=True)
|
||||
+274
@@ -0,0 +1,274 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Picdrop-Zugriff via SFTP: Galerien/Ordner listen, Bilder ziehen & hochladen."""
|
||||
from __future__ import annotations
|
||||
|
||||
import posixpath
|
||||
import stat
|
||||
from contextlib import contextmanager
|
||||
from typing import Iterator
|
||||
|
||||
import paramiko
|
||||
|
||||
from . import config
|
||||
|
||||
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".bmp"}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _sftp() -> Iterator[paramiko.SFTPClient]:
|
||||
transport = paramiko.Transport((config.PICDROP_HOST, config.PICDROP_PORT))
|
||||
transport.banner_timeout = 30
|
||||
transport.connect(username=config.PICDROP_USER, password=config.PICDROP_PASSWORD)
|
||||
try:
|
||||
sftp = paramiko.SFTPClient.from_transport(transport)
|
||||
sftp.get_channel().settimeout(60)
|
||||
yield sftp
|
||||
finally:
|
||||
transport.close()
|
||||
|
||||
|
||||
def list_dir(path: str = "/") -> list[dict]:
|
||||
"""Eintraege eines Verzeichnisses (Name, ob Ordner, Groesse)."""
|
||||
with _sftp() as sftp:
|
||||
out = []
|
||||
for attr in sftp.listdir_attr(path):
|
||||
is_dir = stat.S_ISDIR(attr.st_mode)
|
||||
out.append({
|
||||
"name": attr.filename,
|
||||
"is_dir": is_dir,
|
||||
"size": attr.st_size,
|
||||
"path": posixpath.join(path, attr.filename),
|
||||
})
|
||||
out.sort(key=lambda e: (not e["is_dir"], e["name"].lower()))
|
||||
return out
|
||||
|
||||
|
||||
def list_images(path: str) -> list[dict]:
|
||||
return [
|
||||
e for e in list_dir(path)
|
||||
if not e["is_dir"] and posixpath.splitext(e["name"])[1].lower() in IMAGE_EXTS
|
||||
]
|
||||
|
||||
|
||||
def find_gallery(name_substr: str, root: str = "/", max_depth: int = 3) -> list[str]:
|
||||
"""Sucht Ordner, deren Name den Teilstring enthaelt (case-insensitive)."""
|
||||
needle = name_substr.lower()
|
||||
matches: list[str] = []
|
||||
|
||||
def walk(path: str, depth: int) -> None:
|
||||
if depth > max_depth:
|
||||
return
|
||||
try:
|
||||
entries = list_dir(path)
|
||||
except OSError:
|
||||
return
|
||||
for e in entries:
|
||||
if e["is_dir"]:
|
||||
if needle in e["name"].lower():
|
||||
matches.append(e["path"])
|
||||
walk(e["path"], depth + 1)
|
||||
|
||||
walk(root, 0)
|
||||
return matches
|
||||
|
||||
|
||||
def download(remote_path: str) -> bytes:
|
||||
with _sftp() as sftp:
|
||||
with sftp.open(remote_path, "rb") as fh:
|
||||
fh.prefetch()
|
||||
return fh.read()
|
||||
|
||||
|
||||
def upload(remote_path: str, data: bytes) -> None:
|
||||
with _sftp() as sftp:
|
||||
# Zielordner sicherstellen
|
||||
d = posixpath.dirname(remote_path)
|
||||
parts, cur = d.strip("/").split("/"), ""
|
||||
for p in parts:
|
||||
cur = cur + "/" + p
|
||||
try:
|
||||
sftp.stat(cur)
|
||||
except FileNotFoundError:
|
||||
sftp.mkdir(cur)
|
||||
with sftp.open(remote_path, "wb") as fh:
|
||||
fh.write(data)
|
||||
Reference in New Issue
Block a user