From 6a976790ad7721fb218cf1b8e45a70617d731789 Mon Sep 17 00:00:00 2001 From: Pyrrhic Date: Sun, 20 Nov 2022 17:18:16 -0600 Subject: [PATCH] Use multiprocessing for faces --- Dockerfile | 2 +- Dockerfile-arm32v6 | 2 +- requirements.txt | 1 + src/faces.py | 65 ++++++++++++++++++++++++------------- src/{lib => }/faces_util.py | 45 +++++++++++++------------ 5 files changed, 70 insertions(+), 45 deletions(-) rename src/{lib => }/faces_util.py (82%) diff --git a/Dockerfile b/Dockerfile index 9f5b3b9..8a338ba 100644 --- a/Dockerfile +++ b/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 ./ diff --git a/Dockerfile-arm32v6 b/Dockerfile-arm32v6 index 55dbef4..81231b3 100644 --- a/Dockerfile-arm32v6 +++ b/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 ./ diff --git a/requirements.txt b/requirements.txt index ac2382b..5464907 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ pillow py-cord python-dotenv markovchain +argparse diff --git a/src/faces.py b/src/faces.py index 25ba73e..f96bd4d 100644 --- a/src/faces.py +++ b/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: diff --git a/src/lib/faces_util.py b/src/faces_util.py similarity index 82% rename from src/lib/faces_util.py rename to src/faces_util.py index a1499ee..11113e7 100644 --- a/src/lib/faces_util.py +++ b/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)