Roman commited on
Commit
5369c66
•
1 Parent(s): d9d229c

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

Browse files
app.py CHANGED
@@ -4,18 +4,19 @@ import os
4
  import shutil
5
  import subprocess
6
  import time
7
-
8
  import gradio as gr
9
  import numpy
10
  import requests
 
 
11
  from common import (
12
  AVAILABLE_FILTERS,
13
  CLIENT_TMP_PATH,
 
14
  EXAMPLES,
15
  FILTERS_PATH,
16
  INPUT_SHAPE,
17
  KEYS_PATH,
18
- WRONG_KEYS_PATH,
19
  REPO_DIR,
20
  SERVER_URL,
21
  )
@@ -33,12 +34,23 @@ def decrypt_output_with_wrong_key(encrypted_image, filter_name):
33
  filter_path = FILTERS_PATH / f"{filter_name}/deployment"
34
 
35
  # Instantiate the client interface and generate a new private key
36
- wrong_client = FHEClient(filter_path, WRONG_KEYS_PATH, filter_name)
37
  wrong_client.generate_private_and_evaluation_keys(force=True)
38
 
39
  # Deserialize, decrypt and post-process the encrypted output using the new private key
40
  output_image = wrong_client.deserialize_decrypt_post_process(encrypted_image)
41
 
 
 
 
 
 
 
 
 
 
 
 
42
  return output_image
43
 
44
 
@@ -73,8 +85,8 @@ def get_client(user_id, filter_name):
73
  """
74
  return FHEClient(
75
  FILTERS_PATH / f"{filter_name}/deployment",
76
- KEYS_PATH / f"{filter_name}_{user_id}",
77
  filter_name,
 
78
  )
79
 
80
 
@@ -95,31 +107,32 @@ def get_client_file_path(name, user_id, filter_name):
95
  def clean_temporary_files(n_keys=20):
96
  """Clean keys and encrypted images.
97
 
98
- A maximum of n_keys keys are allowed to be stored. Once this limit is reached, the oldest are
99
- deleted.
100
 
101
  Args:
102
- n_keys (int): The maximum number of keys to be stored. Default to 20.
103
 
104
  """
105
- # Get the oldest files in the key directory
106
- list_files = sorted(KEYS_PATH.iterdir(), key=os.path.getmtime)
107
 
108
  # If more than n_keys keys are found, remove the oldest
109
  user_ids = []
110
- if len(list_files) > n_keys:
111
- n_files_to_delete = len(list_files) - n_keys
112
- for p in list_files[:n_files_to_delete]:
113
- user_ids.append(p.name)
114
- shutil.rmtree(p)
115
 
116
  # Get all the encrypted objects in the temporary folder
117
- list_files_tmp = CLIENT_TMP_PATH.iterdir()
 
118
 
119
- # Delete all files related to the current user
120
- for file in list_files_tmp:
121
  for user_id in user_ids:
122
- if file.name.endswith(f"{user_id}.npy"):
123
  file.unlink()
124
 
125
 
 
4
  import shutil
5
  import subprocess
6
  import time
 
7
  import gradio as gr
8
  import numpy
9
  import requests
10
+ from itertools import chain
11
+
12
  from common import (
13
  AVAILABLE_FILTERS,
14
  CLIENT_TMP_PATH,
15
+ SERVER_TMP_PATH,
16
  EXAMPLES,
17
  FILTERS_PATH,
18
  INPUT_SHAPE,
19
  KEYS_PATH,
 
20
  REPO_DIR,
21
  SERVER_URL,
22
  )
 
34
  filter_path = FILTERS_PATH / f"{filter_name}/deployment"
35
 
36
  # Instantiate the client interface and generate a new private key
37
+ wrong_client = FHEClient(filter_path, filter_name)
38
  wrong_client.generate_private_and_evaluation_keys(force=True)
39
 
40
  # Deserialize, decrypt and post-process the encrypted output using the new private key
41
  output_image = wrong_client.deserialize_decrypt_post_process(encrypted_image)
42
 
43
+ # For filters that are expected to output black and white images, generate two other random
44
+ # channels for better display
45
+ if filter_name in ["black and white", "ridge detection"]:
46
+ # Green channel
47
+ wrong_client.generate_private_and_evaluation_keys(force=True)
48
+ output_image[:, :, 1] = wrong_client.deserialize_decrypt_post_process(encrypted_image)[:, :, 0]
49
+
50
+ # Blue channel
51
+ wrong_client.generate_private_and_evaluation_keys(force=True)
52
+ output_image[:, :, 2] = wrong_client.deserialize_decrypt_post_process(encrypted_image)[:, :, 0]
53
+
54
  return output_image
55
 
56
 
 
85
  """
86
  return FHEClient(
87
  FILTERS_PATH / f"{filter_name}/deployment",
 
88
  filter_name,
89
+ key_dir=KEYS_PATH / f"{filter_name}_{user_id}",
90
  )
91
 
92
 
 
107
  def clean_temporary_files(n_keys=20):
108
  """Clean keys and encrypted images.
109
 
110
+ A maximum of n_keys keys and associated temporary files are allowed to be stored. Once this
111
+ limit is reached, the oldest files are deleted.
112
 
113
  Args:
114
+ n_keys (int): The maximum number of keys and associated files to be stored. Default to 20.
115
 
116
  """
117
+ # Get the oldest key files in the key directory
118
+ key_dirs = sorted(KEYS_PATH.iterdir(), key=os.path.getmtime)
119
 
120
  # If more than n_keys keys are found, remove the oldest
121
  user_ids = []
122
+ if len(key_dirs) > n_keys:
123
+ n_keys_to_delete = len(key_dirs) - n_keys
124
+ for key_dir in key_dirs[:n_keys_to_delete]:
125
+ user_ids.append(key_dir.name)
126
+ shutil.rmtree(key_dir)
127
 
128
  # Get all the encrypted objects in the temporary folder
129
+ client_files = CLIENT_TMP_PATH.iterdir()
130
+ server_files = SERVER_TMP_PATH.iterdir()
131
 
132
+ # Delete all files related to the ids whose keys were deleted
133
+ for file in chain(client_files, server_files):
134
  for user_id in user_ids:
135
+ if user_id in file.name:
136
  file.unlink()
137
 
138
 
client_server_interface.py CHANGED
@@ -55,8 +55,8 @@ class FHEDev:
55
  """Initialize the FHE interface.
56
 
57
  Args:
58
- path_dir (str): The path to the directory where the circuit is saved.
59
  filter (Filter): The filter to use in the FHE interface.
 
60
  """
61
 
62
  self.filter = filter
@@ -83,12 +83,13 @@ class FHEDev:
83
  class FHEClient:
84
  """Client interface to encrypt and decrypt FHE data associated to a Filter."""
85
 
86
- def __init__(self, path_dir, key_dir, filter_name):
87
  """Initialize the FHE interface.
88
 
89
  Args:
90
  path_dir (Path): The path to the directory where the circuit is saved.
91
- key_dir (Path): The path to the directory where the keys are stored.
 
92
  """
93
  self.path_dir = path_dir
94
  self.key_dir = key_dir
 
55
  """Initialize the FHE interface.
56
 
57
  Args:
 
58
  filter (Filter): The filter to use in the FHE interface.
59
+ path_dir (str): The path to the directory where the circuit is saved.
60
  """
61
 
62
  self.filter = filter
 
83
  class FHEClient:
84
  """Client interface to encrypt and decrypt FHE data associated to a Filter."""
85
 
86
+ def __init__(self, path_dir, filter_name, key_dir=None):
87
  """Initialize the FHE interface.
88
 
89
  Args:
90
  path_dir (Path): The path to the directory where the circuit is saved.
91
+ filter_name (str): The filter's name to consider.
92
+ key_dir (Path): The path to the directory where the keys are stored. Default to None.
93
  """
94
  self.path_dir = path_dir
95
  self.key_dir = key_dir
common.py CHANGED
@@ -8,13 +8,11 @@ REPO_DIR = Path(__file__).parent
8
  # This repository's main necessary folders
9
  FILTERS_PATH = REPO_DIR / "filters"
10
  KEYS_PATH = REPO_DIR / ".fhe_keys"
11
- WRONG_KEYS_PATH = REPO_DIR / ".wrong_keys"
12
  CLIENT_TMP_PATH = REPO_DIR / "client_tmp"
13
  SERVER_TMP_PATH = REPO_DIR / "server_tmp"
14
 
15
  # Create the necessary folders
16
  KEYS_PATH.mkdir(exist_ok=True)
17
- WRONG_KEYS_PATH.mkdir(exist_ok=True)
18
  CLIENT_TMP_PATH.mkdir(exist_ok=True)
19
  SERVER_TMP_PATH.mkdir(exist_ok=True)
20
 
 
8
  # This repository's main necessary folders
9
  FILTERS_PATH = REPO_DIR / "filters"
10
  KEYS_PATH = REPO_DIR / ".fhe_keys"
 
11
  CLIENT_TMP_PATH = REPO_DIR / "client_tmp"
12
  SERVER_TMP_PATH = REPO_DIR / "server_tmp"
13
 
14
  # Create the necessary folders
15
  KEYS_PATH.mkdir(exist_ok=True)
 
16
  CLIENT_TMP_PATH.mkdir(exist_ok=True)
17
  SERVER_TMP_PATH.mkdir(exist_ok=True)
18
 
filters/black and white/deployment/client.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:7d870611c586c24a8edde68e7974bc70b5bb7240a6bc6e28ba45bd50c11c8719
3
  size 378
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:009a1f53994ae2776806eef65f516b40b8366b7e16d28b4037bdeebf3f742261
3
  size 378
filters/black and white/deployment/server.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:0113e57bf1016aa8089bfc95ce4eae847e6f69f8d921716718af19672aeffd9a
3
  size 5871
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f1e3227d75f87cad545ba34391f23b1aa37fc808cb6a2e3ec34f76e3ab1af0b9
3
  size 5871
filters/blur/deployment/client.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:0181978730e5cdc7fc669916d949ce2f9a783d867362f731d36388917e3030ea
3
  size 391
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ad80a31b1f0aae22fa3520319f30ccea41210223946e74c00a8c7493cc1c7b46
3
  size 391
filters/blur/deployment/server.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:f9334b5f3b113ddc97fe19a7fdadea7c6f1dbc97123aa1a94470c7b88174a0fb
3
  size 8716
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5c58d04b6c75c926a0d98dc3376761dea78445a27224d63ab37e01a062b1781d
3
  size 8716
filters/identity/deployment/client.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:d8bd5dcf676ff27ea3e272f588fca78ff05db1ec922b5ee018182174d8f144e9
3
  size 376
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:69590288fb2da2e41d772572e959e6ea01e26836ae412e36f927b1a4e7998fb0
3
  size 376
filters/identity/deployment/server.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:62a76ee3f3688b0dab084bf415e2a83b3aafbd4b66f1aec43b3f867818c205e4
3
  size 2537
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c8221f0473d76d0df4242e03b339068276e865b67748cb87f4a3c71eebf99ffe
3
  size 2537
filters/inverted/deployment/client.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:d8bd5dcf676ff27ea3e272f588fca78ff05db1ec922b5ee018182174d8f144e9
3
  size 376
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:433a5406b400f59ada754ebd0ff1cc55cdcffaf628acf82f24820fdb2b3d3fb4
3
  size 376
filters/inverted/deployment/server.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:ddfb48c9ceaeacec0a0ea13edbc3af545621deef9b66ae3a573e9c74295d23bb
3
  size 4152
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:16ee970649f7dcc0be9fe5ca4de1fd0077b307956761aa45b48ae78eb3ab2f7a
3
  size 4152
filters/ridge detection/deployment/client.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:f31154f5c36f5974944129b513bddd33f74a2880d0d47503e0b6d202421d5d50
3
  size 396
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:233f8eb49ac2a5d4fe6bb761c1234249c18f13c6a6158548f5489f78bdb780ed
3
  size 396
filters/ridge detection/deployment/server.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:c5f75c8f54a669c5694c8aa5432c3ef1ff37cf375b7e17ab50e5fb33e50da93c
3
  size 6697
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1b662157b5b7517a18b0a6cacb9c1800b124424f38a5a4230829c7a1ddb2a62b
3
  size 6697
filters/rotate/deployment/client.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:d8bd5dcf676ff27ea3e272f588fca78ff05db1ec922b5ee018182174d8f144e9
3
  size 376
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:433a5406b400f59ada754ebd0ff1cc55cdcffaf628acf82f24820fdb2b3d3fb4
3
  size 376
filters/rotate/deployment/server.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:300962338ed0a7a4d47e1a679a36561e0888d1aafa2dc246630db946b9592313
3
  size 4387
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b55e20b3c46d9609a1f2c0be3b3db5bd0d4e288061a1428e17c9abbe67ed3e10
3
  size 4387
filters/sharpen/deployment/client.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:74f5d14e6f5ccd47560c6a5194986005ba60251d1311a7ea9cf3ff7ee7b794f4
3
  size 396
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:789caedfe972a3108187614b66d15bca0a64fa0bb24840cbec9e8f93211f3750
3
  size 396
filters/sharpen/deployment/server.zip CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:84da9ef7cebe7869be764ba5b5e3c959168722deb6e93bce4ad12c308bb4cad0
3
  size 8735
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:23fe93ee2d0154c15eeae9aacc2f1b4bb41c3a832b9c71b91406e0190138e451
3
  size 8735
server.py CHANGED
@@ -40,8 +40,8 @@ def send_input(
40
  ):
41
  """Send the inputs to the server."""
42
  # Retrieve the encrypted input image and the evaluation key paths
43
- encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
44
- evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)
45
 
46
  # Write the files using the above paths
47
  with encrypted_image_path.open("wb") as encrypted_image, evaluation_key_path.open(
@@ -58,8 +58,8 @@ def run_fhe(
58
  ):
59
  """Execute the filter on the encrypted input image using FHE."""
60
  # Retrieve the encrypted input image and the evaluation key paths
61
- encrypted_image_path = get_server_file_path("encrypted_image", filter, user_id)
62
- evaluation_key_path = get_server_file_path("evaluation_key", filter, user_id)
63
 
64
  # Read the files using the above paths
65
  with encrypted_image_path.open("rb") as encrypted_image_file, evaluation_key_path.open(
@@ -77,7 +77,7 @@ def run_fhe(
77
  fhe_execution_time = round(time.time() - start, 2)
78
 
79
  # Retrieve the encrypted output image path
80
- encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)
81
 
82
  # Write the file using the above path
83
  with encrypted_output_path.open("wb") as encrypted_output:
@@ -93,7 +93,7 @@ def get_output(
93
  ):
94
  """Retrieve the encrypted output image."""
95
  # Retrieve the encrypted output image path
96
- encrypted_output_path = get_server_file_path("encrypted_output", filter, user_id)
97
 
98
  # Read the file using the above path
99
  with encrypted_output_path.open("rb") as encrypted_output_file:
 
40
  ):
41
  """Send the inputs to the server."""
42
  # Retrieve the encrypted input image and the evaluation key paths
43
+ encrypted_image_path = get_server_file_path("encrypted_image", user_id, filter)
44
+ evaluation_key_path = get_server_file_path("evaluation_key", user_id, filter)
45
 
46
  # Write the files using the above paths
47
  with encrypted_image_path.open("wb") as encrypted_image, evaluation_key_path.open(
 
58
  ):
59
  """Execute the filter on the encrypted input image using FHE."""
60
  # Retrieve the encrypted input image and the evaluation key paths
61
+ encrypted_image_path = get_server_file_path("encrypted_image", user_id, filter)
62
+ evaluation_key_path = get_server_file_path("evaluation_key", user_id, filter)
63
 
64
  # Read the files using the above paths
65
  with encrypted_image_path.open("rb") as encrypted_image_file, evaluation_key_path.open(
 
77
  fhe_execution_time = round(time.time() - start, 2)
78
 
79
  # Retrieve the encrypted output image path
80
+ encrypted_output_path = get_server_file_path("encrypted_output", user_id, filter)
81
 
82
  # Write the file using the above path
83
  with encrypted_output_path.open("wb") as encrypted_output:
 
93
  ):
94
  """Retrieve the encrypted output image."""
95
  # Retrieve the encrypted output image path
96
+ encrypted_output_path = get_server_file_path("encrypted_output", user_id, filter)
97
 
98
  # Read the file using the above path
99
  with encrypted_output_path.open("rb") as encrypted_output_file: