memory cache mask

app.py

@@ -8,12 +8,23 @@ import os
 from scipy import ndimage
 # from dotenv import load_dotenv
 import numpy as np
+import hashlib
+import queue
 
 # load_dotenv()
 
 
-openai.api_key = os.getenv('OPENAI_API_KEY')
-removebg_api_key = os.getenv('REMOVEBG_API_KEY')
+cache = dict()
+que = queue.Queue(30)
+
+
+def save_to_memory_cache(key, file):
+    print('save mask')
+    cache[key] = file
+    que.put(key)
+    if que.full():
+        rkey = que.get()
+        del cache[rkey]
 
 
 def get_circle_footprint(size):
@@ -22,7 +33,8 @@ def get_circle_footprint(size):
     return fp
 
 
-def test(prompt, img_path, mask_margin):
+def request_mask(img_path):
+    removebg_api_key = os.getenv('REMOVEBG_API_KEY')
     response = requests.post(
         'https://api.remove.bg/v1.0/removebg',
         files={'image_file': open(img_path, 'rb')},
@@ -34,39 +46,67 @@ def test(prompt, img_path, mask_margin):
     )
     if response.status_code == requests.codes.ok:
         zipFile = zipfile.ZipFile(io.BytesIO(response.content))
-        maskIm = Image.open(io.BytesIO(zipFile.read('alpha.png')))
-
-        alpha= maskIm.getchannel(0)
-        if mask_margin > 0:
-            inflated_alpha = ndimage.maximum_filter(input=np.array(alpha), footprint=get_circle_footprint(mask_margin))
-            alpha = Image.fromarray(np.uint8(inflated_alpha))
-        maskIm.paste((255),[0,0,maskIm.size[0],maskIm.size[1]])
-        maskIm.putalpha(alpha)
-
-        maskFile = io.BytesIO()
-        maskIm.save(maskFile, format='PNG')
-        maskFile.seek(0)
-
-        response = openai.Image.create_edit(
-            image=open(img_path, "rb"),
-            mask=maskFile,
-            prompt=prompt,
-            n=1,
-            size="512x512"
-        )
-        return response['data'][0]['url']
-
-        # with open('no-bg.zip', 'wb') as out:
-        #     out.write(response.content)
+        maskImFile = zipFile.read('alpha.png')
+        return maskImFile
     else:
         print("Error:", response.status_code, response.text)
+        return None
+
+
+def get_file_hash(path):
+    with open(path, 'rb') as inputfile:
+        fh = hashlib.sha256()
+        fb = inputfile.read(65536)
+        while len(fb) > 0:
+            fh.update(fb)
+            fb = inputfile.read(65536)
+    return fh.hexdigest()
+
+
+def process_image(prompt, img_path, mask_margin):
+    openai.api_key = os.getenv('OPENAI_API_KEY')
 
+    hsh = get_file_hash(img_path)
+    print('hash',hsh)
 
+    maskImFile = None
+    if hsh in cache:
+        maskImFile = cache[hsh]
+    else:
+        maskImFile = request_mask(img_path)
+        if maskImFile != None:
+            save_to_memory_cache(hsh, maskImFile)
+        else:
+            print('no mask received')
+            return 'https://i.imgur.com/DUd0OWN.png'
+
+    maskIm = Image.open(io.BytesIO(maskImFile))
+
+    alpha = maskIm.getchannel(0)
+    if mask_margin > 0:
+        inflated_alpha = ndimage.maximum_filter(input=np.array(
+            alpha), footprint=get_circle_footprint(mask_margin))
+        alpha = Image.fromarray(np.uint8(inflated_alpha))
+    maskIm.paste((255), [0, 0, maskIm.size[0], maskIm.size[1]])
+    maskIm.putalpha(alpha)
+
+    maskFile = io.BytesIO()
+    maskIm.save(maskFile, format='PNG')
+    maskFile.seek(0)
+
+    response = openai.Image.create_edit(
+        image=open(img_path, "rb"),
+        mask=maskFile,
+        prompt=prompt,
+        n=1,
+        size="512x512"
+    )
+    return response['data'][0]['url']
 
 # demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo = gr.Interface(test, inputs = [
+demo = gr.Interface(process_image, inputs=[
     'text',
-    gr.Image(type='filepath', shape=(500,500), label='image'),
+    gr.Image(type='filepath', shape=(500, 500), label='image'),
     gr.Slider(minimum=0, maximum=10, value=5, step=1, label="mask margin")
 ], outputs=["image"])