Limit pixel count for pixels

src/config.py.example

@@ -29,3 +29,6 @@ face_dims = [
     [245, 125, 82, 82],
     [98, 131, 87, 87]
 ]
+
+# Max total pixels for a face image to have, set to -1 to disable
+face_max_pixels = 2_000_000

src/faces.py

@@ -26,40 +26,42 @@ async def on_message(message: discord.Message, client, conf):
               if attachment.size > MAX_SIZE:
                 raise TooBigException
               if attachment.proxy_url != '404':
-                await do_face(original_message, 'attachment_picture.png', lambda : attachment.save('attachment_picture.png', use_cached=True))
+                await do_face(original_message, 'attachment_picture.png', conf, lambda : attachment.save('attachment_picture.png', use_cached=True))
                 return
         if len(message.embeds) > 0:
           for embed in message.embeds:
             if embed.type == 'image':
-              await do_face(original_message, 'attachment_picture.png', lambda : download_file(embed.url, 'attachment_picture.png'))
+              await do_face(original_message, 'attachment_picture.png', conf, lambda : download_file(embed.url, 'attachment_picture.png'))
               return
         if message.reference != None:
           message = await message.channel.fetch_message(message.reference.message_id)
           recursions += 1
         else:
           return
   except TooBigException:
-    await message.channel.send('Your image is too big :(')
+    await message.reply('Your image is too big :(')
 
-async def do_face(message: discord.Message, name, do_download):
+async def do_face(message: discord.Message, name, conf, do_download):
   try:
     logging.info(f'{datetime.now(timezone.utc)} Starting face processing for user {message.author.display_name} {message.author.id}')
     await do_download()
     found, path = faces_util.get_face_replace(name)
     if found:
-      await message.channel.send(file=discord.File(path))
+      await message.reply(file=discord.File(path))
     else:
-      await message.channel.send('No face found :(')
+      await message.reply('No face found :(')
     if os.path.exists('face_detected.png'): 
       os.remove('face_detected.png')
     if os.path.exists(name): 
       os.remove(name)
   except discord.NotFound:
-    logging.exception(f'{datetime.now(timezone.utc)} Face on_message')
-    await message.channel.send('Attachment not found :(')
+    logging.exception(f'{datetime.now(timezone.utc)} Face Attachment not found')
+    await message.reply('Attachment not found :(')
   except discord.HTTPException:
-    logging.exception(f'{datetime.now(timezone.utc)} Face on_message')
-    await message.channel.send('Could not download attachment :(')
+    logging.exception(f'{datetime.now(timezone.utc)} Face HTTP error')
+    await message.reply('Could not download attachment :(')
+  except faces_util.TooManyPixelsException as tmp:
+    await message.reply(f'Image has too many pixels. Image: {tmp.pixels:,} Max: {conf.face_max_pixels:,} :(')
 
 async def download_file(url, out_filename):
   async with aiohttp.ClientSession() as session:

src/lib/faces_util.py

@@ -6,12 +6,20 @@
 import config as conf
 from PIL import Image
 
+class TooManyPixelsException(Exception):
+  def __init__(self, pixels: int):
+    self.pixels = pixels
+
 face_cascade = cv2.CascadeClassifier('lib/opencv/face_detector.xml')
 eye_cascade = cv2.CascadeClassifier('lib/opencv/eye_detector.xml')
 smile_cascade = cv2.CascadeClassifier('lib/opencv/smile_detector.xml')
 
-def load_image_rgba(path):
+def load_image_rgba(path, max_pixels = -1):
   og_image = Image.open(path)
+  img_size = og_image.size
+  pixels = img_size[0] * img_size[1]
+  if max_pixels != -1 and pixels > max_pixels:
+    raise TooManyPixelsException(pixels)
   rgba_image = og_image.convert(mode='RGBA')
   img = np.array(rgba_image)
   # B and R channels are swapped and there didn't seem to be a better way to do this
@@ -79,7 +87,7 @@ def replace_faces(gray, frame, replacement, replacement_dims):
 
 def get_face_replace(replacing_image):
   face_replace = load_image_rgba(conf.face_picture)
-  attachment_pic = load_image_rgba(replacing_image)
+  attachment_pic = load_image_rgba(replacing_image, conf.face_max_pixels)
   gray = cv2.cvtColor(attachment_pic, cv2.COLOR_RGBA2GRAY)
   found,img = replace_faces(gray, attachment_pic, face_replace, conf.face_dims)