diff --git a/admin_bot/admin_bot/cogs/image_cog.py b/admin_bot/admin_bot/cogs/image_cog.py index ad29971..4e38599 100644 --- a/admin_bot/admin_bot/cogs/image_cog.py +++ b/admin_bot/admin_bot/cogs/image_cog.py @@ -26,6 +26,7 @@ def register_image_tasks(self): Dynamically register context menu commands for image tasks. """ from utilities.image_utils import IMAGE_TASKS # Import the registered tasks + from utilities.image_utils import CHAT_IMAGE_TASKS for task_name in IMAGE_TASKS: # Create a context menu for each task @@ -38,6 +39,15 @@ def register_image_tasks(self): # Add the command to the bot self.bot.tree.add_command(context_menu) + for task_name in CHAT_IMAGE_TASKS: + command = app_commands.Command( + name=task_name, + callback=self.process_image_chat, + description=task_name.title(), + ) + command.task_name = task_name + self.bot.tree.add_command(command) + async def process_image_context( self, interaction: discord.Interaction, message: discord.Message ): @@ -47,6 +57,36 @@ async def process_image_context( task_name = interaction.command.task_name await self.process_image_task(interaction, task_name, message) + async def process_image_chat( + self, interaction: discord.Interaction, attachment: discord.Attachment + ): + task_name = interaction.command.task_name + + if not attachment.content_type.startswith("image/"): + await interaction.response.send_message( + "The attachment is not an image!", ephemeral=True + ) + return + + await interaction.response.defer() # Acknowledge the interaction + image_bytes = await attachment.read() + + try: + # Process the image + input_image = load_image_from_bytes(image_bytes) + output_image = apply_image_task(input_image, task_name) + output_bytes = save_image_to_bytes(output_image) + + # Send the result + await interaction.followup.send( + file=discord.File(io.BytesIO(output_bytes), f"{task_name}.png") + ) + except Exception as e: + logging.error(f"Error processing image task '{task_name}': {e}") + await interaction.followup.send( + "Failed to process the image.", ephemeral=True + ) + async def process_image_task( self, interaction: discord.Interaction, diff --git a/admin_bot/admin_bot/resources/reaction.png b/admin_bot/admin_bot/resources/reaction.png new file mode 100644 index 0000000..137c509 Binary files /dev/null and b/admin_bot/admin_bot/resources/reaction.png differ diff --git a/admin_bot/admin_bot/utilities/image_tasks.py b/admin_bot/admin_bot/utilities/image_tasks.py index a6ba701..ffe3c68 100644 --- a/admin_bot/admin_bot/utilities/image_tasks.py +++ b/admin_bot/admin_bot/utilities/image_tasks.py @@ -1,5 +1,5 @@ from PIL import Image -from utilities.image_utils import register_image_task +from utilities.image_utils import register_image_task, register_chat_image_task import numpy @@ -171,3 +171,16 @@ def pin_task(input_image: Image.Image) -> Image.Image: result_image.paste(template, (0, 0), template) return result_image + + +# You can only have 5 context menu commands. +# Use @register_chat_image_task to register an +# image command as a regular slash command + + +@register_chat_image_task("live") +def live_task(input_image: Image.Image) -> Image.Image: + template = Image.open("admin_bot/resources/reaction.png").convert("RGBA") + template.paste(input_image.resize((660, 378)), (10, 113)) + template.paste(input_image.resize((223, 84)), (140, 16)) + return template diff --git a/admin_bot/admin_bot/utilities/image_utils.py b/admin_bot/admin_bot/utilities/image_utils.py index 3f88d1e..b21aaab 100644 --- a/admin_bot/admin_bot/utilities/image_utils.py +++ b/admin_bot/admin_bot/utilities/image_utils.py @@ -7,6 +7,9 @@ # Registry for image manipulation tasks IMAGE_TASKS: Dict[str, Callable[[Image.Image], Image.Image]] = {} +# Registry for chat command image manipulation tasks +CHAT_IMAGE_TASKS: Dict[str, Callable[[Image.Image], Image.Image]] = {} + def register_image_task(name: str): """ @@ -22,6 +25,20 @@ def decorator(func: Callable[[Image.Image], Image.Image]): return decorator +def register_chat_image_task(name: str): + """ + Decorator to register an image task. + + :param name: The name of the task. + """ + + def decorator(func: Callable[[Image.Image], Image.Image]): + CHAT_IMAGE_TASKS[name] = func + return func + + return decorator + + def apply_image_task(image: Image.Image, task_name: str) -> Image.Image: """ Apply a registered image task to the given image. @@ -30,10 +47,13 @@ def apply_image_task(image: Image.Image, task_name: str) -> Image.Image: :param task_name: The task to apply. :return: The resulting image. """ - if task_name not in IMAGE_TASKS: - raise ValueError(f"Task '{task_name}' is not registered.") + if task_name in IMAGE_TASKS: + return IMAGE_TASKS[task_name](image) + + if task_name in CHAT_IMAGE_TASKS: + return CHAT_IMAGE_TASKS[task_name](image) - return IMAGE_TASKS[task_name](image) + raise ValueError(f"Task '{task_name}' is not registered.") def load_image_from_bytes(image_bytes: bytes) -> Image.Image: