diff --git a/jarvis/cogs/image.py b/jarvis/cogs/image.py index 27df52e..a9af120 100644 --- a/jarvis/cogs/image.py +++ b/jarvis/cogs/image.py @@ -5,13 +5,16 @@ from io import BytesIO import aiohttp import cv2 import numpy as np -from dis_snek import MessageContext, Scale, Snake, message_command +from dis_snek import InteractionContext, Scale, Snake from dis_snek.models.discord.embed import EmbedField from dis_snek.models.discord.file import File -from dis_snek.models.snek.command import cooldown -from dis_snek.models.snek.cooldowns import Buckets - -from jarvis.utils import build_embed, convert_bytesize, unconvert_bytesize +from dis_snek.models.discord.message import Attachment +from dis_snek.models.snek.application_commands import ( + OptionTypes, + slash_command, + slash_option, +) +from jarvis_core.util import build_embed, convert_bytesize, unconvert_bytesize MIN_ACCURACY = 0.80 @@ -26,39 +29,65 @@ class ImageCog(Scale): def __init__(self, bot: Snake): self.bot = bot self._session = aiohttp.ClientSession() - self.tgt_match = re.compile(r"([0-9]*\.?[0-9]*?) ?([KMGTP]?B)", re.IGNORECASE) + self.tgt_match = re.compile(r"([0-9]*\.?[0-9]*?) ?([KMGTP]?B?)", re.IGNORECASE) def __del__(self): self._session.close() - async def _resize(self, ctx: MessageContext, target: str, url: str = None) -> None: - if not target: - await ctx.send("Missing target size, i.e. 200KB.") + @slash_command(name="resize", description="Resize an image") + @slash_option( + name="target", + description="Target size, i.e. 200KB", + opt_type=OptionTypes.STRING, + required=True, + ) + @slash_option( + name="attachment", + description="Image to resize", + opt_type=OptionTypes.ATTACHMENT, + required=False, + ) + @slash_option( + name="url", + description="URL to download and resize", + opt_type=OptionTypes.STRING, + required=False, + ) + async def _resize( + self, ctx: InteractionContext, target: str, attachment: Attachment = None, url: str = None + ) -> None: + if not attachment and not url: + await ctx.send("A URL or attachment is required", ephemeral=True) + return + + if attachment and not attachment.content_type.startswith("image"): + await ctx.send("Attachment must be an image", ephemeral=True) return tgt = self.tgt_match.match(target) if not tgt: - await ctx.send(f"Invalid target format ({target}). Expected format like 200KB") + await ctx.send( + f"Invalid target format ({target}). Expected format like 200KB", ephemeral=True + ) return tgt_size = unconvert_bytesize(float(tgt.groups()[0]), tgt.groups()[1]) if tgt_size > unconvert_bytesize(8, "MB"): - await ctx.send("Target too large to send. Please make target < 8MB") + await ctx.send("Target too large to send. Please make target < 8MB", ephemeral=True) return - file = None - filename = None - if ctx.message.attachments is not None and len(ctx.message.attachments) > 0: - file = await ctx.message.attachments[0].read() - filename = ctx.message.attachments[0].filename - elif url is not None: - async with self._session.get(url) as resp: - if resp.status == 200: - file = await resp.read() - filename = url.split("/")[-1] - else: - ctx.send("Missing file as either attachment or URL.") - size = len(file) + if attachment: + url = attachment.url + filename = attachment.filename + else: + filename = url.split("/")[-1] + + data = None + async with self._session.get(url) as resp: + if resp.status == 200: + data = await resp.read() + + size = len(data) if size <= tgt_size: await ctx.send("Image already meets target.") return @@ -67,43 +96,38 @@ class ImageCog(Scale): accuracy = 0.0 - while len(file) > tgt_size or (len(file) <= tgt_size and accuracy < MIN_ACCURACY): - old_file = file + while len(data) > tgt_size or (len(data) <= tgt_size and accuracy < MIN_ACCURACY): + old_file = data - buffer = np.frombuffer(file, dtype=np.uint8) + buffer = np.frombuffer(data, dtype=np.uint8) img = cv2.imdecode(buffer, flags=-1) width = int(img.shape[1] * ratio) height = int(img.shape[0] * ratio) new_img = cv2.resize(img, (width, height)) - file = cv2.imencode(".png", new_img)[1].tobytes() - accuracy = (len(file) / tgt_size) * 100 + data = cv2.imencode(".png", new_img)[1].tobytes() + accuracy = (len(data) / tgt_size) * 100 if accuracy <= 0.50: - file = old_file + data = old_file ratio += 0.1 else: - ratio = max(tgt_size / len(file) - 0.02, 0.65) + ratio = max(tgt_size / len(data) - 0.02, 0.65) - bufio = BytesIO(file) - accuracy = (len(file) / tgt_size) * 100 + bufio = BytesIO(data) + accuracy = (len(data) / tgt_size) * 100 fields = [ EmbedField("Original Size", convert_bytesize(size), False), - EmbedField("New Size", convert_bytesize(len(file)), False), + EmbedField("New Size", convert_bytesize(len(data)), False), EmbedField("Accuracy", f"{accuracy:.02f}%", False), ] embed = build_embed(title=filename, description="", fields=fields) embed.set_image(url="attachment://resized.png") await ctx.send( embed=embed, - file=File(file=bufio, filename="resized.png"), + file=File(file=bufio, file_name="resized.png"), ) - @message_command(name="resize") - @cooldown(bucket=Buckets.USER, rate=1, interval=60) - async def _resize_pref(self, ctx: MessageContext, target: str, url: str = None) -> None: - await self._resize(ctx, target, url) - def setup(bot: Snake) -> None: """Add ImageCog to J.A.R.V.I.S."""