telegram_bots/bots/modules/overlay.py

48 lines
2.0 KiB
Python

import os
from tempfile import TemporaryDirectory
from time import sleep
from PIL import Image
from django.db import models
from telegram import Update
from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters, CommandHandler
from bots.models import TelegramBotModuleConfig
class OverlayBotModuleConfig(TelegramBotModuleConfig):
start_text = models.TextField(null=True, blank=True)
type_error_text = models.TextField(null=True, blank=True)
comment = models.TextField(blank=True, null=True)
image = models.ImageField()
MODULE_NAME = 'Overlay'
def start_handler(self, update: Update, ctx: CallbackContext):
if self.start_text:
update.effective_message.reply_text(self.start_text)
def message_handler(self, update: Update, ctx: CallbackContext):
if self.type_error_text:
update.effective_message.reply_text(self.type_error_text)
def photo_handler(self, update: Update, ctx: CallbackContext):
with TemporaryDirectory() as d:
src = os.path.join(d, 'src.jpg')
out = os.path.join(d, 'out.png')
q = update.effective_message.photo[-1].get_file().download(src)
im = Image.open(q).convert('RGBA') # type: Image.Image
overlay = Image.open(self.image.path).convert('RGBA') # type: Image.Image
w, h = im.size
min_side = min(w, h)
overlay = overlay.resize((min_side, min_side), Image.LANCZOS)
im.paste(overlay, ((w - min_side) // 2, (h - min_side) // 2), overlay)
im.save(out)
update.effective_message.reply_photo(open(out, 'rb'), caption=self.comment)
def build_dispatcher(self, dispatcher: Dispatcher):
dispatcher.add_handler(CommandHandler('start', self.start_handler))
dispatcher.add_handler(MessageHandler(Filters.photo, self.photo_handler))
dispatcher.add_handler(MessageHandler(Filters.all, self.message_handler))
return dispatcher