Use multiprocessing for faces

Dockerfile

@@ -6,7 +6,7 @@ WORKDIR /code
 # opencv dependencies https://stackoverflow.com/a/63377623
 RUN apt-get update && apt-get install -y ffmpeg libsm6 libxext6 && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
-RUN pip3 --no-cache-dir install -U py-cord python-dotenv markovchain pillow opencv-python
+RUN pip3 --no-cache-dir install -U py-cord python-dotenv markovchain pillow opencv-python argparse
 
 COPY .env src/ secret-scholars-bot-config.json markov.txt ./
 

Dockerfile-arm32v6

@@ -7,7 +7,7 @@ WORKDIR /code
 
 ENV PYTHONPATH=$PYTHONPATH:/usr/lib/python3.10/site-packages
 
-RUN pip3 --no-cache-dir install -U py-cord python-dotenv markovchain pillow
+RUN pip3 --no-cache-dir install -U py-cord python-dotenv markovchain pillow argparse
 
 COPY .env src/ secret-scholars-bot-config.json markov.txt ./
 

requirements.txt

@@ -4,3 +4,4 @@ pillow
 py-cord
 python-dotenv
 markovchain
+argparse

src/faces.py

@@ -1,10 +1,10 @@
-import lib.faces_util as faces_util
+import faces_util
 import os
 import discord
 import logging
 from datetime import datetime, timezone
 import aiohttp
-from asyncio import run
+import asyncio
 
 BYTES_IN_MEGABYTE = 1048576
 MAX_SIZE = BYTES_IN_MEGABYTE * 8
@@ -13,6 +13,8 @@ class TooBigException(Exception):
 
 allowable_attachment_types = ['image/png', 'image/jpeg']
 
+lock = asyncio.Lock()
+
 async def on_message(message: discord.Message, client, conf):
   try:
     if message.content.startswith("$face"):
@@ -42,26 +44,45 @@ async def on_message(message: discord.Message, client, conf):
     await message.reply('Your image is too big :(')
 
 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.reply(file=discord.File(path))
-    else:
-      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 Attachment not found')
-    await message.reply('Attachment not found :(')
-  except discord.HTTPException:
-    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:,} :(')
+  logging.info('before lock')
+  async with lock:
+    logging.info('in lock')
+    infile = name
+    outfile = 'face_detected.png'
+    try:
+      logging.info(f'{datetime.now(timezone.utc)} Starting face processing for user {message.author.display_name} {message.author.id}')
+      await do_download()
+
+      proc = await asyncio.create_subprocess_shell(
+        f'python3 faces_util.py --infile={infile} --outfile={outfile} --maxpixels={conf.face_max_pixels}',
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE)
+
+      stdout, stderr = await proc.communicate()
+      if proc.returncode == 0:
+        if stdout.decode().rstrip() == 'True':
+          await message.reply(file=discord.File(outfile))
+        else:
+          await message.reply('No face found :(')
+      else:
+        errormessage = stderr.decode().rstrip()
+        if errormessage.startswith('TooManyPixelsException'):
+          pixels = int(errormessage.split(':')[1])
+          await message.reply(f'Image has too many pixels. Image: {pixels:,} Max: {conf.face_max_pixels:,} :(')
+        else:
+          await message.reply('Something went wrong :(')
+          logging.exception(f'{datetime.now(timezone.utc)} Face subprocess error {errormessage}')
+    except discord.NotFound:
+      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 HTTP error')
+      await message.reply('Could not download attachment :(')
+    finally:
+      if os.path.exists(outfile):
+        os.remove(outfile)
+      if os.path.exists(infile):
+        os.remove(infile)
 
 async def download_file(url, out_filename):
   async with aiohttp.ClientSession() as session:

src/lib/faces_util.py → src/faces_util.py

@@ -1,10 +1,10 @@
 import cv2
 import numpy as np
 import math
-import aiohttp
-import os
 import config as conf
 from PIL import Image
+import argparse
+import sys
 
 class TooManyPixelsException(Exception):
   def __init__(self, pixels: int):
@@ -85,33 +85,36 @@ def replace_faces(gray, frame, replacement, replacement_dims):
 
   return found, frame
 
-def get_face_replace(replacing_image):
+def get_face_replace(infile, outfile, max_pixels=conf.face_max_pixels):
   face_replace = load_image_rgba(conf.face_picture)
-  attachment_pic = load_image_rgba(replacing_image, conf.face_max_pixels)
+  attachment_pic = load_image_rgba(infile, max_pixels)
   gray = cv2.cvtColor(attachment_pic, cv2.COLOR_RGBA2GRAY)
   found,img = replace_faces(gray, attachment_pic, face_replace, conf.face_dims)
 
   if found:
-    cv2.imwrite('face_detected.png', img)
-    return True, 'face_detected.png'
+    cv2.imwrite(outfile, img)
+    return True
 
-  return False, None
+  return False
 
 if __name__ == '__main__':
-  img = load_image_rgba('student_small.png')
-  # img = load_image_rgba('group.png')
-  student = np.array(img)
-
-  face = load_image_rgba('face.png')
-  face_dims = [
-    [43, 208, 430, 430],
-    [245, 125, 82, 82],
-    [98, 131, 87, 87]
-  ]
-  gray = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY)
-  found,img = replace_faces(gray, img, face, face_dims)
-  cv2.imwrite("face_detected.png", img)
-  print('Successfully saved')
+  parser = argparse.ArgumentParser(description='Apply faces to faces')
+  parser.add_argument('--infile', default='student_small.png', help='The file to apply faces to')
+  parser.add_argument('--outfile', default='face_detected.png', help='The file to save the modified image to')
+  parser.add_argument('--maxpixels', type=int, default=-1, help='Max total pixels for the infile to have, set to -1 to disable (default: -1)')
+  args = parser.parse_args()
+
+  # face = load_image_rgba('face.png')
+  # face_dims = [
+  #   [43, 208, 430, 430],
+  #   [245, 125, 82, 82],
+  #   [98, 131, 87, 87]
+  # ]
+  try:
+    print(get_face_replace(args.infile, args.outfile, args.maxpixels))
+  except TooManyPixelsException as err:
+    print(f'TooManyPixelsException:{err.pixels}', file=sys.stderr)
+    exit(1)
 
 # A video with faces applied for testing purposes
 # video_capture = cv2.VideoCapture(0)