Add image cog

jarvis/cogs/admin.py

@@ -3,6 +3,12 @@ from discord.ext import commands
 
 
 class AdminCog(commands.Cog):
+    """
+    Guild admin functions
+
+    Used to manage guilds
+    """
+
     def __init__(self, bot):
         self.bot = bot
 

jarvis/cogs/image.py

@@ -0,0 +1,74 @@
+import jarvis
+from discord import File
+import cv2
+import aiohttp
+import re
+import numpy as np
+from jarvis.utils import convert_bytesize, unconvert_bytesize
+from discord.ext import commands
+from io import BytesIO
+
+
+class ImageCog(commands.Cog):
+    """
+    Image processing functions for J.A.R.V.I.S.
+
+    May be categorized under util later
+    """
+
+    def __init__(self, bot):
+        self.bot = bot
+        self._session = aiohttp.ClientSession()
+        self.tgt_match = re.compile(
+            r"([0-9]*\.?[0-9]*?) ?([KMGTP]?B)", re.IGNORECASE
+        )
+
+    @commands.command(name="resize", help="Resize an image")
+    async def _resize(self, ctx, target: str, url: str = None):
+        if not target:
+            await ctx.send("Missing target size, i.e. 200KB.")
+            return
+
+        tgt = self.tgt_match.match(target)
+        if not tgt:
+            await ctx.send(
+                "Invalid target format ({}).".format(target)
+                + " Expected format like 200KB"
+            )
+            return
+
+        tgt_size = unconvert_bytesize(*tgt.groups())
+        file = None
+        if (
+            ctx.message.attachments is not None
+            and len(ctx.message.attachments) > 0
+        ):
+            file = await ctx.message.attachments[0].read()
+        elif url is not None:
+            async with self._session.get(url) as resp:
+                if resp.status == 200:
+                    file = await resp.read()
+        else:
+            ctx.send("Missing file as either attachment or URL.")
+
+        size = len(file)
+        if size <= tgt_size:
+            await ctx.send("Image already meets target.")
+            return
+
+        ratio = size / tgt_size
+
+        buffer = np.frombuffer(file, 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))
+        buffer = cv2.imencode(".png", new_img)[1].tobytes()
+        bufio = BytesIO(buffer)
+        await ctx.send(file=File(bufio))
+
+
+def setup(bot):
+    bot.add_cog(ImageCog(bot))

jarvis/cogs/jokes.py

@@ -10,6 +10,12 @@ from datetime import datetime
 
 
 class JokeCog(commands.Cog):
+    """
+    Joke library for J.A.R.V.I.S.
+
+    May adapt over time to create jokes using machine learning
+    """
+
     def __init__(self, bot):
         self.bot = bot
         self.db = db.create_connection()

jarvis/cogs/owner.py

@@ -6,6 +6,12 @@ from jarvis.utils import update
 
 
 class OwnerCog(commands.Cog):
+    """
+    J.A.R.V.I.S. management cog
+
+    Used by admins to control core J.A.R.V.I.S. systems
+    """
+
     def __init__(self, bot):
         self.bot = bot
         self.admins = get_config().admins

jarvis/cogs/util.py

@@ -6,6 +6,12 @@ from discord.ext import commands
 
 
 class UtilCog(commands.Cog):
+    """
+    Utility functions for J.A.R.V.I.S.
+
+    Mostly system utility functions, but may change over time
+    """
+
     def __init__(self, bot):
         self.bot = bot
         self.config = config.get_config()

jarvis/utils/__init__.py

@@ -10,12 +10,21 @@ __all__ = ["field", "db"]
 
 def convert_bytesize(bytes: int) -> str:
     bytes = float(bytes)
-    sizes = ["B", "KB", "MB", "GB"]
+    sizes = ["B", "KB", "MB", "GB", "TB", "PB"]
     size = 0
     while bytes >= 1024 and size < len(sizes) - 1:
         bytes = bytes / 1024
         size += 1
-    return f"{bytes:0.2f}{sizes[size]}"
+    return "{:0.3f} {}".format(bytes, sizes[size])
+
+
+def unconvert_bytesize(size, ending: str):
+    ending = ending.upper()
+    sizes = ["B", "KB", "MB", "GB", "TB", "PB"]
+    if ending == "B":
+        return size
+    # Rounding is only because bytes cannot be partial
+    return round(size * (1024 ** sizes.index(ending)))
 
 
 def get_prefix(bot, message):