156 lines
5 KiB
Python
156 lines
5 KiB
Python
"""J.A.R.V.I.S. image processing cog."""
|
|
import logging
|
|
import re
|
|
from io import BytesIO
|
|
|
|
import aiohttp
|
|
import cv2
|
|
import numpy as np
|
|
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.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
|
|
|
|
|
|
class ImageCog(Scale):
|
|
"""
|
|
Image processing functions for J.A.R.V.I.S.
|
|
|
|
May be categorized under util later
|
|
"""
|
|
|
|
def __init__(self, bot: Snake):
|
|
self.bot = bot
|
|
self.logger = logging.getLogger(__name__)
|
|
self._session = aiohttp.ClientSession()
|
|
self.tgt_match = re.compile(r"([0-9]*\.?[0-9]*?) ?([KMGTP]?B?)", re.IGNORECASE)
|
|
|
|
def __del__(self):
|
|
self._session.close()
|
|
|
|
@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:
|
|
await ctx.defer()
|
|
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", ephemeral=True
|
|
)
|
|
return
|
|
|
|
try:
|
|
tgt_size = unconvert_bytesize(float(tgt.groups()[0]), tgt.groups()[1])
|
|
except ValueError:
|
|
await ctx.send("Failed to read your target size. Try a more sane one", ephemeral=True)
|
|
return
|
|
|
|
if tgt_size > unconvert_bytesize(8, "MB"):
|
|
await ctx.send("Target too large to send. Please make target < 8MB", ephemeral=True)
|
|
return
|
|
elif tgt_size < 1024:
|
|
await ctx.send("Sizes < 1KB are extremely unreliable and are disabled", ephemeral=True)
|
|
return
|
|
|
|
if attachment:
|
|
url = attachment.url
|
|
filename = attachment.filename
|
|
else:
|
|
filename = url.split("/")[-1]
|
|
|
|
data = None
|
|
try:
|
|
async with self._session.get(url) as resp:
|
|
resp.raise_for_status()
|
|
if resp.content_type in ["image/jpeg", "image/png"]:
|
|
data = await resp.read()
|
|
else:
|
|
await ctx.send(
|
|
"Unsupported content type. Please send a URL to a JPEG or PNG",
|
|
ephemeral=True,
|
|
)
|
|
return
|
|
except Exception:
|
|
await ctx.send("Failed to retrieve image. Please verify url", ephemeral=True)
|
|
return
|
|
|
|
size = len(data)
|
|
if size <= tgt_size:
|
|
await ctx.send("Image already meets target.", ephemeral=True)
|
|
return
|
|
|
|
ratio = max(tgt_size / size - 0.02, 0.50)
|
|
|
|
accuracy = 0.0
|
|
|
|
while len(data) > tgt_size or (len(data) <= tgt_size and accuracy < MIN_ACCURACY):
|
|
old_file = data
|
|
|
|
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))
|
|
data = cv2.imencode(".png", new_img)[1].tobytes()
|
|
accuracy = (len(data) / tgt_size) * 100
|
|
if accuracy <= 0.50:
|
|
data = old_file
|
|
ratio += 0.1
|
|
else:
|
|
ratio = max(tgt_size / len(data) - 0.02, 0.65)
|
|
|
|
bufio = BytesIO(data)
|
|
accuracy = (len(data) / tgt_size) * 100
|
|
fields = [
|
|
EmbedField("Original Size", convert_bytesize(size), 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, file_name="resized.png"),
|
|
)
|
|
|
|
|
|
def setup(bot: Snake) -> None:
|
|
"""Add ImageCog to J.A.R.V.I.S."""
|
|
ImageCog(bot)
|