diff --git a/code/bot.py b/code/bot.py index 74ee4b2..4b7ed19 100644 --- a/code/bot.py +++ b/code/bot.py @@ -1,18 +1,18 @@ -import asyncio, logging, warnings +import asyncio +import logging -from init import * +from init import bot, dp from handlers import router async def main() -> None: dp.include_router(router) await bot.delete_webhook(drop_pending_updates=True) - await dp.start_polling(bot, allowed_updates=dp.resolve_used_update_types(), skipUpdates=True) + await dp.start_polling(bot, allowed_updates=dp.resolve_used_update_types()) if __name__ == "__main__": logging.basicConfig(level=logging.INFO) - # warnings.filterwarnings("ignore") try: asyncio.run(main()) except KeyboardInterrupt: pass diff --git a/code/config.py b/code/config.py index 31d1e5a..c21a46a 100644 --- a/code/config.py +++ b/code/config.py @@ -1,4 +1,6 @@ -botToken = "XXXXXXXXXXXXXXXXXXXXXXXX" # @circlechek_bot +import os + +botToken = os.getenv("BOT_TOKEN", "XXXXXXXXXXXXXXXXXXXXXXXX") # @circlechek_bot WARN = ("Предупреждение! Кружочек может быть с некорректными метаданными, в этом случае программа выдаст ошибку " "либо произведёт видео с браком! Обязательно перепроверяйте получившиеся видео!\n\n") diff --git a/code/handlers.py b/code/handlers.py index 9108e25..0cccee0 100644 --- a/code/handlers.py +++ b/code/handlers.py @@ -1,10 +1,15 @@ -from aiogram import types, F, Router +import asyncio +import logging +import os +import shutil +import tempfile + +from aiogram import F, Router +from aiogram.exceptions import TelegramBadRequest from aiogram.types import Message, CallbackQuery, ContentType, InlineKeyboardButton, InlineKeyboardMarkup, FSInputFile from aiogram.filters import Command -import shutil - -from scripts import * +from scripts import Movie, process_frames from config import WARN router = Router() @@ -15,7 +20,7 @@ router = Router() async def helps(msg: Message) -> None: buttons = [[InlineKeyboardButton(text="ФУНКЦИИ", callback_data="fun"), InlineKeyboardButton(text="АВТОР", callback_data="auth")]] - keyboard = InlineKeyboardMarkup(inline_keyboard=buttons, row_width=2) + keyboard = InlineKeyboardMarkup(inline_keyboard=buttons) await msg.answer(text="Кружочичек — бот для обработки видео и кружочков в Телеграме. Для начала работы " "вам достаточно скинуть квадратное видео не дольше минуты в чат, чтобы получить кружок. " @@ -33,7 +38,7 @@ async def author(call: CallbackQuery) -> None: 'двух предложенных вариантах.' '\n\nЯ же пишу подобные небольшие проекты, о которых вы можете узнать ' - 'больше на моём GitHub.') + 'больше на моём GitHub.') @router.callback_query(F.data == "fun") @@ -45,78 +50,84 @@ async def function(call: CallbackQuery) -> None: # ОБРАБОТЧИК ВИДЕО @router.message(F.content_type == ContentType.VIDEO) async def video_to_circle(msg: Message) -> None: - video = "../data/circles/" + str(msg.chat.id) + ".mp4" + video = f"../data/circles/{msg.chat.id}_{msg.message_id}.mp4" await msg.reply("Началась обработка видео! Оно должно быть квадратным и не дольше одной минуты, в ином " "случае бот в ответ вернёт вам изначальное видео, а не кружочек!") - await msg.bot.download(file=msg.video.file_id, destination=video) - await msg.answer_video_note(video_note=FSInputFile(video)) - os.remove(video) + try: + await msg.bot.download(file=msg.video.file_id, destination=video) + await msg.answer_video_note(video_note=FSInputFile(video)) + except TelegramBadRequest as err: + logging.warning("Не удалось сделать кружочек: %s", err) + await msg.answer("Не получилось сделать кружочек. Видео должно быть квадратным, не дольше минуты " + "и не больше 20 МБ.") + finally: + if os.path.exists(video): + os.remove(video) # ОБРАБОТЧИК КРУЖОЧКОВ @router.message(F.content_type == ContentType.VIDEO_NOTE) async def video_note(msg: Message) -> None: - buttons = [[InlineKeyboardButton(text="Градиент", callback_data="grad"), - InlineKeyboardButton(text="Блюр", callback_data="blur")]] - keyboard = types.InlineKeyboardMarkup(inline_keyboard=buttons, row_width=2) + # номер сообщения в кнопках: если прислать несколько кружочков подряд, каждый обработается свой + buttons = [[InlineKeyboardButton(text="Градиент", callback_data=f"grad:{msg.message_id}"), + InlineKeyboardButton(text="Блюр", callback_data=f"blur:{msg.message_id}")]] + keyboard = InlineKeyboardMarkup(inline_keyboard=buttons) msg_answer = await msg.reply("Кружочек загружается...") - await msg.bot.download(file=msg.video_note.file_id, destination="../data/video_notes/" + str(msg.chat.id) + ".mp4") + await msg.bot.download(file=msg.video_note.file_id, + destination=f"../data/video_notes/{msg.chat.id}_{msg.message_id}.mp4") await msg_answer.edit_text(text=WARN + "Какой тип фона в углах вы выберите?", reply_markup=keyboard) -@router.callback_query(lambda call: call.data == "grad" or call.data == "blur") +@router.callback_query(F.data.regexp(r"^(grad|blur):\d+$")) async def work_part(call: CallbackQuery) -> None: - msg = await call.message.edit_text(WARN + "Начало обработки!") - - video_name = str(msg.chat.id) + mode, note_id = call.data.split(":") + video_name = f"{call.message.chat.id}_{note_id}" video_file = video_name + ".mp4" + source = "../data/video_notes/" + video_file - path = f"../data/videos/{video_name}" - os.mkdir(path) - - path_video = path + "/" + video_file - os.replace("../data/video_notes/" + video_file, path_video) - - path_frames = path + f"/frames" - os.mkdir(path_frames) - - path_background = path + f"/background" - os.mkdir(path_background) - - msg = await msg.edit_text(WARN + "Этап: 1 - Обработка видео.") - video = Movie(path_video, video_name, path) - procces = video.split_into_frames() - - if not procces: - await msg.edit_text("Возникла ошибка! Кружочек невозможно обработать!") - shutil.rmtree(path) + if not os.path.exists(source): + # кнопку нажали повторно или кружочек уже обработан + await call.answer("Этот кружочек уже обработан. Пришлите его ещё раз.", show_alert=True) return - msg = await msg.edit_text(WARN + "Этап: 2 - Обработка кадров.") - frames = list(sorted(os.listdir(path_frames))) + msg = await call.message.edit_text(WARN + "Начало обработки!") - for frame in frames: - img = Frame(path_frames + "/" + frame, path_background + "/" + f'{frame[:-5]}-g.jpeg') - width, height = img.size() + # отдельная папка на каждую обработку и уборка в любом случае — иначе после первой же ошибки + # повторная обработка в этом чате падала на уже существующей папке + path = tempfile.mkdtemp(prefix=video_name + "-", dir="../data/videos") + try: + path_video = path + "/" + video_file + os.replace(source, path_video) - if call.data == "blur": - img.blur(width, height) - else: - r, g, b = img.medium_color(); nearly = 10 - img.gradient(width, height, (r - nearly, g - nearly, b - nearly), - (r + nearly, g + nearly, b + nearly), (True, False, False)) + path_frames = path + "/frames" + os.mkdir(path_frames) - img.unity_image() + path_background = path + "/background" + os.mkdir(path_background) - msg = await msg.edit_text(WARN + "Этап: 3 - Объединение кадров.") - video.unity_into_video(frames) + # тяжёлая обработка — в отдельном потоке, чтобы бот не замирал для остальных пользователей + msg = await msg.edit_text(WARN + "Этап: 1 - Обработка видео.") + video = Movie(path_video, video_name, path) + if not await asyncio.to_thread(video.split_into_frames): + await msg.edit_text("Возникла ошибка! Кружочек невозможно обработать!") + return - msg_final = await msg.edit_text(WARN + "Готово! Видео отправляется...") + msg = await msg.edit_text(WARN + "Этап: 2 - Обработка кадров.") + frames = await asyncio.to_thread(process_frames, path_frames, path_background, mode) - await msg.answer_video(video=FSInputFile(path + "/acs-" + video_file), caption=WARN, - reply_to_message_id=call.message.reply_to_message.message_id) - await msg_final.delete() + msg = await msg.edit_text(WARN + "Этап: 3 - Объединение кадров.") + result = await asyncio.to_thread(video.unity_into_video, frames) - shutil.rmtree(path) + msg_final = await msg.edit_text(WARN + "Готово! Видео отправляется...") + + reply_to = call.message.reply_to_message + await msg.answer_video(video=FSInputFile(result), caption=WARN, + reply_to_message_id=reply_to.message_id if reply_to else None) + await msg_final.delete() + except Exception: + logging.exception("Ошибка обработки кружочка") + await msg.edit_text("Возникла ошибка! Кружочек невозможно обработать!") + finally: + shutil.rmtree(path, ignore_errors=True) diff --git a/code/init.py b/code/init.py index 5154444..278b58d 100644 --- a/code/init.py +++ b/code/init.py @@ -7,26 +7,3 @@ from config import botToken bot = Bot(token=botToken, default=DefaultBotProperties(parse_mode=ParseMode.HTML)) dp = Dispatcher(storage=MemoryStorage()) - - -import logging, functools, sys - -# logging.basicConfig(level=logging.INFO, filename='../data/info.log',filemode="w", format="%(asctime)s %(levelname)s %(message)s") - -def bug_report(func): - @functools.wraps(func) - - def wrapper(*args, **kwargs): - - try: - result = func(*args, **kwargs) - logging.info(f"Successful result") - return result - except Exception as err: - logging.error("Ошибка", exc_info=True) - """with open('../data/log.txt', "a") as f: - f.write(repr(e))""" - sys.exit() - - return wrapper - diff --git a/code/scripts.py b/code/scripts.py index 6c21744..90ed2af 100644 --- a/code/scripts.py +++ b/code/scripts.py @@ -1,8 +1,9 @@ -from PIL import Image, ImageDraw, ImageFilter -from moviepy import * -import numpy, os +import logging +import os -from init import bug_report +import numpy +from moviepy import AudioFileClip, ImageSequenceClip, VideoFileClip +from PIL import Image, ImageDraw, ImageFilter class Movie: @@ -11,48 +12,52 @@ class Movie: self.video_name = video_name self.path = path - self.path_frames = "" + self.path_frames = self.path + "/frames" + self.path_audio = self.path + f"/{self.video_name}-audio.mp3" + self.has_audio = False def split_into_frames(self): - video_clip = VideoFileClip(self.video_file) - video_clip.audio.write_audiofile(self.path + f"/{self.video_name}-audio.mp3", logger=None) + try: + video_clip = VideoFileClip(self.video_file) + except Exception: + logging.exception("Не удалось открыть кружочек") + return False - self.path_frames = self.path + "/frames" - step = 1 / 30.0 if video_clip.fps > 60.0 else 1 / video_clip.fps + try: + # у кружочка может не быть звуковой дорожки + if video_clip.audio is not None: + video_clip.audio.write_audiofile(self.path_audio, logger=None) + self.has_audio = True - count = 0 - for current_duration in numpy.arange(0, video_clip.duration, step): - count += 1 + step = 1 / 30.0 if video_clip.fps > 60.0 else 1 / video_clip.fps - str_c = str(count) - num = "00000"[:(len(str_c) * -1)] + str_c + for count, current_duration in enumerate(numpy.arange(0, video_clip.duration, step), start=1): + frame_filename = os.path.join(self.path_frames, f"frame-{count:05}.jpeg") + video_clip.save_frame(frame_filename, current_duration) + except Exception: + logging.exception("Не удалось разобрать кружочек на кадры") + return False + finally: + video_clip.close() - frame_filename = os.path.join(self.path_frames, f"frame-{num}.jpeg") - try: video_clip.save_frame(frame_filename, current_duration) - except Exception: return False - - video_clip.close() return True - @bug_report def unity_into_video(self, frames): video_clip = VideoFileClip(self.video_file) fps = 30.0 if video_clip.fps > 60.0 else video_clip.fps video_clip.close() - os.remove(self.video_file) - clip = ImageSequenceClip(list(map(lambda x: self.path_frames + "/" + x, frames)), fps=fps) - clip.write_videofile(self.video_file, logger=None) + clip = ImageSequenceClip([self.path_frames + "/" + x for x in frames], fps=fps) + if self.has_audio: + audio_clip = AudioFileClip(self.path_audio) + # звук не длиннее видео, иначе последний кадр «замирает» + clip = clip.with_audio(audio_clip.subclipped(0, min(audio_clip.duration, clip.duration))) + + result = self.path + "/acs-" + self.video_name + ".mp4" + clip.write_videofile(result, codec="libx264", audio_codec="aac", logger=None) clip.close() - video_clip = VideoFileClip(self.video_file) - audio_clip = AudioFileClip(self.path + f"/{self.video_name}-audio.mp3") - - video_clip_with_audio = video_clip.with_audio(audio_clip) - video_clip_with_audio.write_videofile(self.path + "/acs-" + self.video_name + ".mp4", logger=None) - - video_clip.close() - audio_clip.close() + return result class Frame: @@ -61,31 +66,14 @@ class Frame: self.background = background def size(self): - image = Image.open(self.filename) - width, height = image.size - image.close() - - return width, height + with Image.open(self.filename) as image: + return image.size def medium_color(self): - img = Image.open(self.filename) - width, height = img.size + with Image.open(self.filename) as img: + pixels = numpy.asarray(img.convert("RGB"), dtype=numpy.float64) - pixels = img.load() - r, g, b = [], [], [] - - for y in range(width): - for x in range(height): - p = pixels[x, y] - r.append(p[0]) - g.append(p[1]) - b.append(p[2]) - - r = sum(r) // len(r) - g = sum(g) // len(g) - b = sum(b) // len(b) - - img.close() + r, g, b = (int(c) for c in pixels.reshape(-1, 3).mean(axis=0)) return r, g, b def gradient(self, width, height, start_list, stop_list, is_horizontal_list): @@ -99,11 +87,12 @@ class Frame: for i, (start, stop, is_horizontal) in enumerate(zip(start_list, stop_list, is_horizontal_list)): result[:, :, i] = get_gradient(start, stop, width, height, is_horizontal) - Image.fromarray(numpy.uint8(result)).save(self.background, quality=95) + # без обрезки тёмные и светлые цвета «переворачиваются» при переводе в uint8 + Image.fromarray(numpy.uint8(numpy.clip(result, 0, 255))).save(self.background, quality=95) def blur(self, width, height): - image = Image.open(self.filename) - image = image.filter(ImageFilter.GaussianBlur(6)) + with Image.open(self.filename) as image: + image = image.filter(ImageFilter.GaussianBlur(6)) part_width = width // 64 part_height = height // 64 @@ -129,3 +118,22 @@ class Frame: im1.close() im2.close() mask.close() + + +def process_frames(path_frames, path_background, mode): + frames = sorted(os.listdir(path_frames)) + + for frame in frames: + img = Frame(path_frames + "/" + frame, path_background + "/" + f"{frame[:-5]}-g.jpeg") + width, height = img.size() + + if mode == "blur": + img.blur(width, height) + else: + r, g, b = img.medium_color(); nearly = 10 + img.gradient(width, height, (r - nearly, g - nearly, b - nearly), + (r + nearly, g + nearly, b + nearly), (True, False, False)) + + img.unity_image() + + return frames diff --git a/requirements.txt b/requirements.txt index 77601cd..d527860 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,14 +1,4 @@ -# Core bot framework -aiogram>=3.0.0 - -# Image processing -Pillow>=9.0.0 - -# Video processing -moviepy>=1.0.3 - -# Mathematical operations -numpy>=1.21.0 - -# Optional: For better video codec support -imageio-ffmpeg>=0.4.0 +aiogram==3.23.0 +moviepy==2.2.1 +Pillow==10.3.0 +numpy==2.2.6