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:
+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()
|
||||
Reference in New Issue
Block a user