chore: Better display for b&w images and fix tmp files cleaning

app.py

@@ -4,18 +4,19 @@ import os
 import shutil
 import subprocess
 import time
-
 import gradio as gr
 import numpy
 import requests
+from itertools import chain
+
 from common import (
     AVAILABLE_FILTERS,
     CLIENT_TMP_PATH,
+    SERVER_TMP_PATH,
     EXAMPLES,
     FILTERS_PATH,
     INPUT_SHAPE,
     KEYS_PATH,
-    WRONG_KEYS_PATH,
     REPO_DIR,
     SERVER_URL,
 )
@@ -33,12 +34,23 @@ def decrypt_output_with_wrong_key(encrypted_image, filter_name):
     filter_path = FILTERS_PATH / f"{filter_name}/deployment"
 
     # Instantiate the client interface and generate a new private key
-    wrong_client = FHEClient(filter_path, WRONG_KEYS_PATH, filter_name)
+    wrong_client = FHEClient(filter_path, filter_name)
     wrong_client.generate_private_and_evaluation_keys(force=True)
 
     # Deserialize, decrypt and post-process the encrypted output using the new private key
     output_image = wrong_client.deserialize_decrypt_post_process(encrypted_image)
 
+    # For filters that are expected to output black and white images, generate two other random 
+    # channels for better display 
+    if filter_name in ["black and white", "ridge detection"]:
+        # Green channel
+        wrong_client.generate_private_and_evaluation_keys(force=True)
+        output_image[:, :, 1] = wrong_client.deserialize_decrypt_post_process(encrypted_image)[:, :, 0]
+
+        # Blue channel
+        wrong_client.generate_private_and_evaluation_keys(force=True)
+        output_image[:, :, 2] = wrong_client.deserialize_decrypt_post_process(encrypted_image)[:, :, 0]
+
     return output_image
 
 
@@ -73,8 +85,8 @@ def get_client(user_id, filter_name):
     """
     return FHEClient(
         FILTERS_PATH / f"{filter_name}/deployment", 
-        KEYS_PATH / f"{filter_name}_{user_id}", 
         filter_name,
+        key_dir=KEYS_PATH / f"{filter_name}_{user_id}", 
     )
 
 
@@ -95,31 +107,32 @@ def get_client_file_path(name, user_id, filter_name):
 def clean_temporary_files(n_keys=20):
     """Clean keys and encrypted images.
 
-    A maximum of n_keys keys are allowed to be stored. Once this limit is reached, the oldest are
-    deleted.
+    A maximum of n_keys keys and associated temporary files are allowed to be stored. Once this 
+    limit is reached, the oldest files are deleted.
 
     Args:
-        n_keys (int): The maximum number of keys to be stored. Default to 20.
+        n_keys (int): The maximum number of keys and associated files to be stored. Default to 20.
 
     """
-    # Get the oldest files in the key directory
-    list_files = sorted(KEYS_PATH.iterdir(), key=os.path.getmtime)
+    # Get the oldest key files in the key directory
+    key_dirs = sorted(KEYS_PATH.iterdir(), key=os.path.getmtime)
 
     # If more than n_keys keys are found, remove the oldest
     user_ids = []
-    if len(list_files) > n_keys:
-        n_files_to_delete = len(list_files) - n_keys
-        for p in list_files[:n_files_to_delete]:
-            user_ids.append(p.name)
-            shutil.rmtree(p)
+    if len(key_dirs) > n_keys:
+        n_keys_to_delete = len(key_dirs) - n_keys
+        for key_dir in key_dirs[:n_keys_to_delete]:
+            user_ids.append(key_dir.name)
+            shutil.rmtree(key_dir)
 
     # Get all the encrypted objects in the temporary folder
-    list_files_tmp = CLIENT_TMP_PATH.iterdir()
+    client_files = CLIENT_TMP_PATH.iterdir()
+    server_files = SERVER_TMP_PATH.iterdir()
 
-    # Delete all files related to the current user
-    for file in list_files_tmp:
+    # Delete all files related to the ids whose keys were deleted
+    for file in chain(client_files, server_files):
         for user_id in user_ids:
-            if file.name.endswith(f"{user_id}.npy"):
+            if user_id in file.name:
                 file.unlink()
 
 

client_server_interface.py

@@ -55,8 +55,8 @@ class FHEDev:
         """Initialize the FHE interface.
 
         Args:
-            path_dir (str): The path to the directory where the circuit is saved.
             filter (Filter): The filter to use in the FHE interface.
+            path_dir (str): The path to the directory where the circuit is saved.
         """
 
         self.filter = filter
@@ -83,12 +83,13 @@ class FHEDev:
 class FHEClient:
     """Client interface to encrypt and decrypt FHE data associated to a Filter."""
 
-    def __init__(self, path_dir, key_dir, filter_name):
+    def __init__(self, path_dir, filter_name, key_dir=None):
         """Initialize the FHE interface.
 
         Args:
             path_dir (Path): The path to the directory where the circuit is saved.
-            key_dir (Path): The path to the directory where the keys are stored.
+            filter_name (str): The filter's name to consider.
+            key_dir (Path): The path to the directory where the keys are stored. Default to None.
         """
         self.path_dir = path_dir
         self.key_dir = key_dir

common.py

@@ -8,13 +8,11 @@ REPO_DIR = Path(__file__).parent
 # This repository's main necessary folders
 FILTERS_PATH = REPO_DIR / "filters"
 KEYS_PATH = REPO_DIR / ".fhe_keys"
-WRONG_KEYS_PATH = REPO_DIR / ".wrong_keys"
 CLIENT_TMP_PATH = REPO_DIR / "client_tmp"
 SERVER_TMP_PATH = REPO_DIR / "server_tmp"
 
 # Create the necessary folders
 KEYS_PATH.mkdir(exist_ok=True)
-WRONG_KEYS_PATH.mkdir(exist_ok=True)
 CLIENT_TMP_PATH.mkdir(exist_ok=True)
 SERVER_TMP_PATH.mkdir(exist_ok=True)
 

filters/black and white/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7d870611c586c24a8edde68e7974bc70b5bb7240a6bc6e28ba45bd50c11c8719
+oid sha256:009a1f53994ae2776806eef65f516b40b8366b7e16d28b4037bdeebf3f742261
 size 378

filters/black and white/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0113e57bf1016aa8089bfc95ce4eae847e6f69f8d921716718af19672aeffd9a
+oid sha256:f1e3227d75f87cad545ba34391f23b1aa37fc808cb6a2e3ec34f76e3ab1af0b9
 size 5871

filters/blur/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0181978730e5cdc7fc669916d949ce2f9a783d867362f731d36388917e3030ea
+oid sha256:ad80a31b1f0aae22fa3520319f30ccea41210223946e74c00a8c7493cc1c7b46
 size 391

filters/blur/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f9334b5f3b113ddc97fe19a7fdadea7c6f1dbc97123aa1a94470c7b88174a0fb
+oid sha256:5c58d04b6c75c926a0d98dc3376761dea78445a27224d63ab37e01a062b1781d
 size 8716

filters/identity/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d8bd5dcf676ff27ea3e272f588fca78ff05db1ec922b5ee018182174d8f144e9
+oid sha256:69590288fb2da2e41d772572e959e6ea01e26836ae412e36f927b1a4e7998fb0
 size 376

filters/identity/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:62a76ee3f3688b0dab084bf415e2a83b3aafbd4b66f1aec43b3f867818c205e4
+oid sha256:c8221f0473d76d0df4242e03b339068276e865b67748cb87f4a3c71eebf99ffe
 size 2537

filters/inverted/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d8bd5dcf676ff27ea3e272f588fca78ff05db1ec922b5ee018182174d8f144e9
+oid sha256:433a5406b400f59ada754ebd0ff1cc55cdcffaf628acf82f24820fdb2b3d3fb4
 size 376

filters/inverted/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ddfb48c9ceaeacec0a0ea13edbc3af545621deef9b66ae3a573e9c74295d23bb
+oid sha256:16ee970649f7dcc0be9fe5ca4de1fd0077b307956761aa45b48ae78eb3ab2f7a
 size 4152

filters/ridge detection/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f31154f5c36f5974944129b513bddd33f74a2880d0d47503e0b6d202421d5d50
+oid sha256:233f8eb49ac2a5d4fe6bb761c1234249c18f13c6a6158548f5489f78bdb780ed
 size 396

filters/ridge detection/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c5f75c8f54a669c5694c8aa5432c3ef1ff37cf375b7e17ab50e5fb33e50da93c
+oid sha256:1b662157b5b7517a18b0a6cacb9c1800b124424f38a5a4230829c7a1ddb2a62b
 size 6697

filters/rotate/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d8bd5dcf676ff27ea3e272f588fca78ff05db1ec922b5ee018182174d8f144e9
+oid sha256:433a5406b400f59ada754ebd0ff1cc55cdcffaf628acf82f24820fdb2b3d3fb4
 size 376

filters/rotate/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:300962338ed0a7a4d47e1a679a36561e0888d1aafa2dc246630db946b9592313
+oid sha256:b55e20b3c46d9609a1f2c0be3b3db5bd0d4e288061a1428e17c9abbe67ed3e10
 size 4387

filters/sharpen/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:74f5d14e6f5ccd47560c6a5194986005ba60251d1311a7ea9cf3ff7ee7b794f4
+oid sha256:789caedfe972a3108187614b66d15bca0a64fa0bb24840cbec9e8f93211f3750
 size 396

filters/sharpen/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:84da9ef7cebe7869be764ba5b5e3c959168722deb6e93bce4ad12c308bb4cad0
+oid sha256:23fe93ee2d0154c15eeae9aacc2f1b4bb41c3a832b9c71b91406e0190138e451
 size 8735

server.py

@@ -40,8 +40,8 @@ def send_input(
 ):
     """Send the inputs to the server."""
     # Retrieve the encrypted input image and the evaluation key paths
-    encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
-    evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)
+    encrypted_image_path = get_server_file_path("encrypted_image", user_id, filter)
+    evaluation_key_path = get_server_file_path("evaluation_key", user_id, filter)
     
     # Write the files using the above paths
     with encrypted_image_path.open("wb") as encrypted_image, evaluation_key_path.open(
@@ -58,8 +58,8 @@ def run_fhe(
 ):
     """Execute the filter on the encrypted input image using FHE."""
     # Retrieve the encrypted input image and the evaluation key paths
-    encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
-    evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)
+    encrypted_image_path = get_server_file_path("encrypted_image", user_id, filter)
+    evaluation_key_path = get_server_file_path("evaluation_key", user_id, filter)
 
     # Read the files using the above paths
     with encrypted_image_path.open("rb") as encrypted_image_file, evaluation_key_path.open(
@@ -77,7 +77,7 @@ def run_fhe(
     fhe_execution_time = round(time.time() - start, 2)
 
     # Retrieve the encrypted output image path
-    encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)
+    encrypted_output_path = get_server_file_path("encrypted_output", user_id, filter)
 
     # Write the file using the above path
     with encrypted_output_path.open("wb") as encrypted_output:
@@ -93,7 +93,7 @@ def get_output(
 ):
     """Retrieve the encrypted output image."""
     # Retrieve the encrypted output image path
-    encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)
+    encrypted_output_path = get_server_file_path("encrypted_output", user_id, filter)
 
     # Read the file using the above path
     with encrypted_output_path.open("rb") as encrypted_output_file: