jarvis/jarvis-bot
jarvis/jarvis-bot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Migrate image cog

Browse source
This commit is contained in:
Zeva Rose 2022-02-21 01:26:07 -07:00
parent 0175fce442
commit b12b109ad5
1 changed files with 64 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

104
jarvis/cogs/image.py
View file

@ -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."""