Zem.TelegramBot/extensions/watermark.py

import os, pathlib, asyncio
from time import time

try:
	from telegram import ForceReply, Update
	from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandler, filters, ConversationHandler
except ModuleNotFoundError:
	Update = None
	class ct:
		DEFAULT_TYPE = None
	ContextTypes = ct()
	print("Cannot import telegram futures")

filter_tables = {
	"v1":{#default
		"resize_filter":"[logo][0]scale2ref=w=oh*mdar:h=ih*1.0[logo][0];",
		"insert_filter":"[0][logo]overlay=(W-w)/2:(H-h)/2:format=auto,format=yuv420p",
		"filter":"[1]format=rgba,colorchannelmixer=aa={transparency}[logo];{resize_filter}{insert_filter}"
	},
	"v2":{#photo
		"filter":"[1]format=rgba,colorchannelmixer=aa={transparency}[logo];[logo][0]scale2ref=oh*mdar:ih*0.2[logo][video];[video][logo]overlay=(main_w-overlay_w):(main_h-overlay_h)"
	},
	"v2.move":{#good
		"insert_filter":"[video][logo]overlay=x='if(lt(mod(t\,16)\,8)\,W-w-W*10/100\,W*10/100)':y='if(lt(mod(t+4\,16)\,8)\,H-h-H*5/100\,H*5/100)'",
		"filter":"[1]colorchannelmixer=aa={transparency}[logo];[logo][0]scale2ref=oh*mdar:ih*0.2[logo][video];{insert_filter}"
	},
	"v2.dvd":{
		"insert_filter":"[video][logo]overlay=x='if(eq(mod(floor(st(8,t*{speed_x})/(W-w)),2),0),mod(ld(8),W-w),W-w-mod(ld(8),W-w))':y='if(eq(mod(floor(st(9,t*{speed_y})/(H-h)),2),0),mod(ld(9),H-h),H-h-mod(ld(9),H-h))'",
		"filter":"[1]colorchannelmixer=aa={transparency}[logo];[logo][0]scale2ref=oh*mdar:ih*0.2[logo][video];{insert_filter}"
	}
}

async def add_watermark(input_file, watermark, transparency, out_ext, filter_version = "v1", speed=(160, 70)):
	resize_filter = filter_tables[filter_version].get("resize_filter", "")
	insert_filter = filter_tables[filter_version].get("insert_filter", "").format(speed_x = speed[0], speed_y = speed[1])
	filter = filter_tables[filter_version]["filter"].format(transparency = transparency, resize_filter = resize_filter, insert_filter = insert_filter)
	cmd = f'ffmpeg -y -i {input_file} -i {watermark} -filter_complex "{filter}" {"-c:v libx264 -c:a copy" if out_ext == "mp4" else ""} {input_file}.{out_ext}'
	print(cmd)
	process = await asyncio.create_subprocess_shell(cmd)
	await process.wait()
	if os.path.exists(f"{input_file}.{out_ext}"):
		return f"{input_file}.{out_ext}"
	return ""

class Extension:
	cmds = ["/watermark"]
	wait_content = {}
	content_directory = os.getenv("CONTENT_DIR") if os.getenv("CONTENT_DIR", None) else "./content"
	png_path = os.getenv("WATERMARK")
	default_transparency=20.0
	watermark_users = []

	def __init__(self, core) -> None:
		self.watermark_users = core.admins
		print(f"Content save path: {self.content_directory}")
		print(f"Watermark users: {' '.join([str(i) for i in self.watermark_users])}")

	async def __call__(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
		if update.message.text and update.message.text.split()[0] in self.cmds:
			return await self.watermark(update, context)
		return await self.content_handler(update, context)

	async def watermark(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
		if not update.message.from_user.id in self.watermark_users:
			return await update.message.reply_text("Это не для таких как ты сделано, друг...")

		if update.message.from_user.id in self.wait_content:
			del self.wait_content[update.message.from_user.id]
			return await update.message.reply_text("Понял понял, не бей, отменяем думанье...")
		else:
			try:
				self.wait_content[update.message.from_user.id] = float(update.message.text.split()[1])
			except:
				self.wait_content[update.message.from_user.id] = self.default_transparency
			return await update.message.reply_text(f"Отлично, жду от тебя ОДНУ фотку/видео/переслал\nНепрозрачность: {self.wait_content[update.message.from_user.id]}%")

	async def content_handler(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
		if not update.message.from_user.id in self.wait_content:
			return
		if update.message.video:
			return await self.video_handler(update, context)
		elif update.message.photo:
			return await self.photo_handler(update, context)

	#Что ниже над будет засунуть в одно, ахуй
	async def photo_handler(self, update:Update, context: ContextTypes.DEFAULT_TYPE):
		if update.message.from_user.id in self.wait_content:
			transparency = self.wait_content[update.message.from_user.id] / 100
			del self.wait_content[update.message.from_user.id]
		else:
			return

		_id = update.message.photo[-1].file_id
		photo = await context.bot.get_file(_id)
		downloaded_photo = await photo.download_to_drive(f"{self.content_directory}/{time()}")
		res = await add_watermark(downloaded_photo.absolute(), self.png_path, transparency, "jpg", "v2")
		if res:
			await update.message.reply_photo(photo=res)
			os.remove(res)
		else:
			await update.message.reply_text("Ошибка при создании фотки")
		if os.path.exists(downloaded_photo.absolute()):
			os.remove(downloaded_photo.absolute())
		return

	async def video_handler(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
		if update.message.from_user.id in self.wait_content:
			transparency = self.wait_content[update.message.from_user.id] / 100
			del self.wait_content[update.message.from_user.id]
		else:
			return

		video = await update.message.video.get_file()
		downloaded_video = await video.download_to_drive(f"{self.content_directory}/{time()}")
		res = await add_watermark(downloaded_video.absolute(), self.png_path, transparency, "mp4", "v2.dvd")
		if res:
			await update.message.reply_video(video=res)
			os.remove(res)
		else:
			await update.message.reply_text("Ошибка при создании видео")
		if os.path.exists(downloaded_video.absolute()):
			os.remove(downloaded_video.absolute())
		return

if __name__ == "__main__":
	import asyncio
	import argparse
	parser = argparse.ArgumentParser()
	parser.add_argument("input", type=str)
	parser.add_argument("watermark", type=str)
	parser.add_argument("--transparency", type=float, default=100.0)
	parser.add_argument("--insert-version", type=str, default="v1")
	parser.add_argument("--out-ext", type=str, default="mp4")
	args = parser.parse_args()

	async def run():
		await add_watermark(args.input, args.watermark, args.transparency / 100, args.out_ext, args.insert_version)

	asyncio.get_event_loop().run_until_complete(run())