Телеграм-бот для супер-разрешения изображений на базе Real-ESRGAN x4 (Python, aiogram)
Представляем телеграм-бота для супер-разрешения изображений на базе Real-ESRGAN ×4. Решение написано на Python с использованием aiogram и поддерживает как GPU (CUDA, FP16), так и CPU. Бот принимает изображения как «фото» или как «документ» (рекомендуется для максимального качества, без сжатия Telegram) и возвращает увеличенную в 4 раза версию. Поддерживаются форматы PNG/JPG/JPEG/TIFF/WEBP, предусмотрен лимит размера 25 МБ.
Технические особенности.
Модель RealESRGAN_x4plus загружается автоматически; веса кешируются на диск.
Явная инициализация RRDBNet и совместимость с realesrgan==0.3.0.
Тайлинг (tile/pad) для устойчивой обработки больших изображений.
Очередь задач через asyncio.Semaphore (по одному изображению за раз) для безопасной работы на GPU.
Конвертация RGB↔BGR (PIL↔NumPy), сохранение результата с качеством JPEG 95.
Хеш-основанные имена файлов вывода и подробный логгинг.
Рекомендована передача токена через переменную окружения TELEGRAM_BOT_TOKEN.
Проект предназначен для быстрого развёртывания сервиса «апскейлинга по запросу» в клинических, исследовательских и медиа-потоках, обеспечивая воспроизводимость, простоту интеграции и высокое качество результата.
# bot_ru.py
import asyncio
import os
import io
import hashlib
import logging
import warnings
from pathlib import Path
import numpy as np
from aiogram import Bot, Dispatcher, F
from aiogram.types import Message, FSInputFile
from aiogram.enums import ParseMode
from aiogram.filters import CommandStart, Command
from PIL import Image
import torch
import requests
from tqdm import tqdm
# ---- Логирование / предупреждения ----
warnings.filterwarnings("ignore", category=UserWarning)
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s:%(name)s:%(message)s",
datefmt="%H:%M:%S",
)
# ---- Конфиг ----
# РЕКОМЕНДАЦИЯ: храните токен в переменной окружения TELEGRAM_BOT_TOKEN
TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "PASTE_YOUR_TOKEN_HERE") # установите перед запуском
bot = Bot(token=TOKEN)
dp = Dispatcher()
MODELS_DIR = Path("models")
OUTPUT_DIR = Path("output")
MODELS_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
# Real-ESRGAN x4 (универсальная модель для фото)
MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
MODEL_PATH = MODELS_DIR / "RealESRGAN_x4plus.pth"
UPSCALE = 4
MAX_BYTES = 25 * 1024 * 1024 # лимит 25 МБ для больших документов
SEM = asyncio.Semaphore(1) # обрабатываем по 1 изображению (безопасно для GPU)
_model = None # кэш модели
# ---- Real-ESRGAN (PyPI 0.3.0) ----
from realesrgan import RealESRGANer
from basicsr.archs.rrdbnet_arch import RRDBNet
def _download_with_progress(url: str, dest: Path):
"""Скачивание файла с прогресс-баром."""
with requests.get(url, stream=True, timeout=120) as r:
r.raise_for_status()
total = int(r.headers.get("content-length", 0))
with open(dest, "wb") as f, tqdm(
total=total, unit="B", unit_scale=True, desc=f"Скачивание {dest.name}"
) as bar:
for chunk in r.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
bar.update(len(chunk))
def _ensure_model_file() -> Path:
"""Гарантирует наличие весов модели на диске."""
if not MODEL_PATH.exists():
MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
_download_with_progress(MODEL_URL, MODEL_PATH)
return MODEL_PATH
def _get_device():
"""Определение устройства (CUDA/CPU)."""
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
def _load_realesrgan():
"""Надёжная загрузка для realesrgan==0.3.0 (нужен явный RRDBNet)."""
global _model
if _model is not None:
return _model
weights_path = _ensure_model_file()
device = _get_device()
# Стандартная архитектура RealESRGAN_x4plus
net = RRDBNet(
num_in_ch=3, num_out_ch=3,
num_feat=64, num_block=23, num_grow_ch=32,
scale=UPSCALE
)
# Загружаем state_dict из чекпоинта (ключи могут различаться)
ckpt = torch.load(str(weights_path), map_location=device)
if isinstance(ckpt, dict):
state = ckpt.get("params_ema") or ckpt.get("params") or ckpt.get("state_dict") or ckpt
else:
state = ckpt
if not isinstance(state, dict):
raise RuntimeError("Некорректный чекпоинт для RealESRGAN_x4plus.pth")
net.load_state_dict(state, strict=False)
# Создаём RealESRGANer
_model = RealESRGANer(
scale=UPSCALE,
model_path=str(weights_path), # обязательно в 0.3.0
model=net, # наш RRDBNet
dni_weight=None,
device=device,
tile=200, tile_pad=10, pre_pad=0, # тайлинг для больших изображений
half=torch.cuda.is_available() # FP16 на GPU
)
logging.info("Модель загружена на устройство: %s", device)
return _model
def _safe_image_open(data: bytes) -> Image.Image:
"""Открытие изображения с переводом в RGB (если нужно)."""
img = Image.open(io.BytesIO(data))
return img.convert("RGB") if img.mode != "RGB" else img
async def _process_and_send(message: Message, photo_bytes: bytes, suffix: str):
"""Апскейл изображения и отправка результата пользователю."""
async with SEM:
model = _load_realesrgan()
img = _safe_image_open(photo_bytes) # PIL (RGB)
# PIL (RGB) -> NumPy (BGR) для RealESRGANer
img_np = np.array(img)[:, :, ::-1]
# Апскейл (RealESRGANer возвращает NumPy BGR)
up_np, _ = model.enhance(img_np, outscale=UPSCALE)
# NumPy (BGR) -> PIL (RGB)
up_img = Image.fromarray(up_np[:, :, ::-1])
# Сохранение + отправка
sha = hashlib.sha1(photo_bytes).hexdigest()[:10]
out_path = OUTPUT_DIR / f"upscaled_x{UPSCALE}_{sha}{suffix.lower()}"
save_kwargs = {"quality": 95} if suffix.lower() in [".jpg", ".jpeg"] else {}
up_img.save(out_path, **save_kwargs)
await message.reply_document(
FSInputFile(out_path),
caption=f"✅ Увеличено Real-ESRGAN x{UPSCALE}",
parse_mode=ParseMode.HTML
)
# ---- Хэндлеры ----
@dp.message(CommandStart())
async def start(message: Message):
await message.answer(
"Привет! Отправь мне изображение, и я увеличу его с помощью Real-ESRGAN x4.\n"
"• Для максимального качества отправляй как Документ (Telegram не сжимает).\n"
"• Поддерживаются PNG/JPG/JPEG/TIFF/WEBP.",
parse_mode=ParseMode.HTML
)
@dp.message(Command("help"))
async def help_cmd(message: Message):
await message.answer(
"Отправь изображение (фото или документ). Я верну увеличенную версию x4 через Real-ESRGAN.",
parse_mode=ParseMode.HTML
)
@dp.message(F.photo)
async def handle_photo(message: Message):
try:
await message.reply("🔧 Обрабатываю изображение… (Real-ESRGAN x4)")
file = await bot.get_file(message.photo[-1].file_id)
f = await bot.download_file(file.file_path)
data = f.read()
await _process_and_send(message, data, suffix=".jpg")
except Exception as e:
logging.exception("photo error")
await message.reply(f"❌ Ошибка: {e}")
@dp.message(F.document)
async def handle_document(message: Message):
try:
mime = (message.document.mime_type or "").lower()
if message.document.file_size and message.document.file_size > MAX_BYTES:
await message.reply("Файл слишком большой (>25 МБ). Пожалуйста, отправьте изображение меньшего размера.")
return
# Принимаем только изображения
if not any(x in mime for x in ["image/", "png", "jpeg", "jpg", "tiff", "bmp", "webp"]):
await message.reply("Пожалуйста, отправьте файл-изображение (PNG/JPG/TIFF/WEBP).", parse_mode=ParseMode.HTML)
return
await message.reply("🔧 Обрабатываю изображение… (Real-ESRGAN x4)")
file = await bot.get_file(message.document.file_id)
f = await bot.download_file(file.file_path)
data = f.read()
# Определяем расширение для сохранения
ext = ".jpg"
if "png" in mime: ext = ".png"
elif "webp" in mime: ext = ".webp"
elif "tif" in mime or "tiff" in mime: ext = ".tif"
elif "bmp" in mime: ext = ".bmp"
await _process_and_send(message, data, suffix=ext)
except Exception as e:
logging.exception("document error")
await message.reply(f"❌ Ошибка: {e}")
async def main():
# Префетч модели (необязательно, просто заранее скачает веса при старте)
try:
_ensure_model_file()
except Exception as e:
logging.warning("Префетч не удался (продолжаем работу): %s", e)
await dp.start_polling(bot, allowed_updates=dp.resolve_used_update_types())
if __name__ == "__main__":
try:
asyncio.run(main())
except (KeyboardInterrupt, SystemExit):
print("Остановлено.")
Комментариев нет:
Отправить комментарий