chore: display encrypted ouput, add sizes and resize input examples

.gitignore

@@ -2,6 +2,7 @@
 .playground/
 .artifacts
 .fhe_keys
+.wrong_keys
 server_tmp/
 client_tmp/
 .artifacts

README.md

@@ -36,6 +36,7 @@ Then, install required packages:
 <!--pytest-codeblocks:skip-->
 
 ```bash
+pip3 install pip --upgrade
 pip3 install -U pip wheel setuptools --ignore-installed
 pip3 install -r requirements.txt --ignore-installed
 ```

app.py

@@ -15,16 +15,28 @@ from common import (
     FILTERS_PATH,
     INPUT_SHAPE,
     KEYS_PATH,
+    WRONG_KEYS_PATH,
     REPO_DIR,
     SERVER_URL,
 )
-from custom_client_server import CustomFHEClient
+from custom_client_server import CustomFHEClient, CustomFHEServer
 
 # Uncomment here to have both the server and client in the same terminal
 subprocess.Popen(["uvicorn", "server:app"], cwd=REPO_DIR)
 time.sleep(3)
 
 
+def decrypt_output_with_wrong_key(encrypted_image, image_filter):
+    filter_path = FILTERS_PATH / f"{image_filter}/deployment"
+
+    wrong_client = CustomFHEClient(filter_path, WRONG_KEYS_PATH)
+    wrong_client.generate_private_and_evaluation_keys(force=True)
+
+    output_image = wrong_client.deserialize_decrypt_post_process(encrypted_image)
+
+    return output_image
+
+
 def shorten_bytes_object(bytes_object, limit=500):
     """Shorten the input bytes object to a given length.
 
@@ -121,8 +133,6 @@ def keygen(image_filter):
     user_id = numpy.random.randint(0, 2**32)
 
     # Retrieve the client API
-    # Currently, the key generation needs to be done after choosing a filter
-    # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2258
     client = get_client(user_id, image_filter)
 
     # Generate a private key
@@ -133,6 +143,10 @@ def keygen(image_filter):
     # execution
     evaluation_key = client.get_serialized_evaluation_keys()
 
+    # Compute the private key's size in Kilobytes
+    private_key_path = next(client.key_dir.iterdir()) / "0_0/secretKey_big" 
+    private_key_size = private_key_path.stat().st_size / 1000
+
     # Save evaluation_key as bytes in a file as it is too large to pass through regular Gradio
     # buttons (see https://github.com/gradio-app/gradio/issues/1877)
     evaluation_key_path = get_client_file_path("evaluation_key", user_id, image_filter)
@@ -140,7 +154,7 @@ def keygen(image_filter):
     with evaluation_key_path.open("wb") as evaluation_key_file:
         evaluation_key_file.write(evaluation_key)
 
-    return (user_id, True)
+    return (user_id, True, private_key_size)
 
 
 def encrypt(user_id, input_image, image_filter):
@@ -159,12 +173,18 @@ def encrypt(user_id, input_image, image_filter):
     if user_id == "":
         raise gr.Error("Please generate the private key first.")
 
+    if input_image is None:
+        raise gr.Error("Please choose an image first.")
+
     # Retrieve the client API
     client = get_client(user_id, image_filter)
 
     # Pre-process, encrypt and serialize the image
     encrypted_image = client.pre_process_encrypt_serialize(input_image)
 
+    # Compute the input's size in Megabytes
+    encrypted_input_size = len(encrypted_image) / 1000000
+
     # Save encrypted_image to bytes in a file, since too large to pass through regular Gradio
     # buttons, https://github.com/gradio-app/gradio/issues/1877
     encrypted_image_path = get_client_file_path("encrypted_image", user_id, image_filter)
@@ -175,7 +195,7 @@ def encrypt(user_id, input_image, image_filter):
     # Create a truncated version of the encrypted image for display
     encrypted_image_short = shorten_bytes_object(encrypted_image)
 
-    return (input_image, encrypted_image_short)
+    return (input_image, encrypted_image_short, encrypted_input_size)
 
 
 def send_input(user_id, image_filter):
@@ -264,17 +284,22 @@ def get_output(user_id, image_filter):
         data=data,
     ) as response:
         if response.ok:
+            encrypted_output = response.content
+
+            # Compute the output's size in Megabytes
+            encrypted_output_size = len(encrypted_output) / 1000000
+
             # Save the encrypted output to bytes in a file as it is too large to pass through regular
             # Gradio buttons (see https://github.com/gradio-app/gradio/issues/1877)
             encrypted_output_path = get_client_file_path("encrypted_output", user_id, image_filter)
 
             with encrypted_output_path.open("wb") as encrypted_output_file:
-                encrypted_output_file.write(response.content)
+                encrypted_output_file.write(encrypted_output)
 
-            # Create a truncated version of the encrypted output for display
-            encrypted_output_image_short = shorten_bytes_object(response.content)
+            # Decrypt the image using a different (wrong) key for display
+            output_image_representation = decrypt_output_with_wrong_key(encrypted_output, image_filter)
 
-            return encrypted_output_image_short
+            return output_image_representation, encrypted_output_size
         else:
             raise gr.Error("Please wait for the FHE execution to be completed.")
 
@@ -328,9 +353,11 @@ with demo:
     )
 
     gr.Markdown("## Client side")
+    gr.Markdown("### Step 1. Upload an image. ")
     gr.Markdown(
-        f"Step 1. Upload an image. It will automatically be resized to shape ({INPUT_SHAPE[0]}x{INPUT_SHAPE[1]})."
-        "The image is however displayed using its original resolution."
+        f"The image will automatically be resized to shape ({INPUT_SHAPE[0]}x{INPUT_SHAPE[1]})."
+        "The image displayed here is however using its original resolution. The true image used "
+        "in this demo can be seen in Step. 8."
     )
     with gr.Row():
         input_image = gr.Image(
@@ -341,33 +368,46 @@ with demo:
             examples=EXAMPLES, inputs=[input_image], examples_per_page=5, label="Examples to use."
         )
 
-    gr.Markdown("Step 2. Choose your filter")
+    gr.Markdown("### Step 2. Choose your filter")
     image_filter = gr.Dropdown(
         choices=AVAILABLE_FILTERS, value="inverted", label="Choose your filter", interactive=True
     )
 
-    gr.Markdown("### Notes")
+    gr.Markdown("#### Notes")
     gr.Markdown(
         """
         - The private key is used to encrypt and decrypt the data and shall never be shared.
         - No public key are required for these filter operators.
         """
     )
-
+    
+    gr.Markdown("### Step 3. Generate the private key.")
     with gr.Row():
-        keygen_button = gr.Button("Step 3. Generate the private key.")
+        keygen_button = gr.Button("Generate the private key.")
 
         keygen_checkbox = gr.Checkbox(label="Private key generated:", interactive=False)
 
-    with gr.Row():
-        encrypt_button = gr.Button("Step 4. Encrypt the image using FHE.")
+        private_key_size = gr.Number(
+            label="Private key size (in kB):", value=0, precision=1, interactive=False
+        )
 
-        user_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
+    user_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
+    
+    gr.Markdown("### Step 4. Encrypt the image using FHE.")
+    with gr.Row():
+        encrypt_button = gr.Button("Encrypt the image using FHE.")
+
+        # FIXME: Display an image representation here
+        # Currently not really possible as Concrete-Numpy does not give the ability to decrypt
+        # an encrypted input 
+        # https://github.com/zama-ai/concrete-ml-internal/issues/2265
+        # https://github.com/zama-ai/concrete-numpy-internal/issues/1753
+        encrypted_input = gr.Textbox(
+            label="Encrypted input representation:", max_lines=2, interactive=False
+        )
 
-        # Display an image representation
-        # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2265
-        encrypted_image = gr.Textbox(
-            label="Encrypted image representation:", max_lines=2, interactive=False
+        encrypted_input_size = gr.Number(
+            label="Encrypted input size (in MB):", value=0, precision=1, interactive=False
         )
 
     gr.Markdown("## Server side")
@@ -377,25 +417,36 @@ with demo:
         "the encrypted results to the client."
     )
 
+    gr.Markdown("### Step 5. Send the encrypted image to the server.")
     with gr.Row():
-        send_input_button = gr.Button("Step 5. Send the encrypted image to the server.")
+        send_input_button = gr.Button("Send the encrypted image to the server.")
 
         send_input_checkbox = gr.Checkbox(label="Encrypted image sent.", interactive=False)
 
+    gr.Markdown("### Step 6. Run FHE execution.")
     with gr.Row():
-        execute_fhe_button = gr.Button("Step 6. Run FHE execution")
+        execute_fhe_button = gr.Button("Run FHE execution.")
 
         fhe_execution_time = gr.Textbox(
-            label="Total FHE execution time (in seconds).", max_lines=1, interactive=False
+            label="Total FHE execution time (in seconds):", max_lines=1, interactive=False
         )
 
+    gr.Markdown("### Step 7. Receive the encrypted output image from the server.")
+    gr.Markdown(
+        "The image displayed here is the encrypted result sent by the server which has been "
+        "decrypted using a different private key. This is only used to visually represent an "
+        "encrypted image."
+    )
     with gr.Row():
-        get_output_button = gr.Button("Step 7. Receive the encrypted output image from the server.")
+        get_output_button = gr.Button("Receive the encrypted output image from the server.")
 
-        # Display an image representation
-        # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2265
-        encrypted_output_image = gr.Textbox(
-            label="Encrypted output image representation:", max_lines=2, interactive=False
+        encrypted_output_representation = gr.Image(
+            label=f"Encrypted output representation ({INPUT_SHAPE[0]}x{INPUT_SHAPE[1]}):", interactive=False
+        )
+        encrypted_output_representation.style(height=256, width=256)
+
+        encrypted_output_size = gr.Number(
+            label="Encrypted output size (in MB):", value=0, precision=1, interactive=False
         )
 
     gr.Markdown("## Client side")
@@ -404,7 +455,12 @@ with demo:
         "private key. Only the client is aware of the original image and its transformed version."
     )
 
-    decrypt_button = gr.Button("Step 8. Decrypt the output")
+    gr.Markdown("### Step 8. Decrypt the output")
+    gr.Markdown(
+        "The image displayed on the left is the input image used during the demo. The output image "
+        "can be seen on the right."
+    )   
+    decrypt_button = gr.Button("Decrypt the output")
 
     # Final input vs output display
     with gr.Row():
@@ -424,14 +480,14 @@ with demo:
     keygen_button.click(
         keygen,
         inputs=[image_filter],
-        outputs=[user_id, keygen_checkbox],
+        outputs=[user_id, keygen_checkbox, private_key_size],
     )
 
     # Button to encrypt inputs on the client side
     encrypt_button.click(
         encrypt,
         inputs=[user_id, input_image, image_filter],
-        outputs=[original_image, encrypted_image],
+        outputs=[original_image, encrypted_input, encrypted_input_size],
     )
 
     # Button to send the encodings to the server using post method
@@ -444,7 +500,9 @@ with demo:
 
     # Button to send the encodings to the server using post method
     get_output_button.click(
-        get_output, inputs=[user_id, image_filter], outputs=[encrypted_output_image]
+        get_output, 
+        inputs=[user_id, image_filter], 
+        outputs=[encrypted_output_representation, encrypted_output_size]
     )
 
     # Button to decrypt the output on the client side

common.py

@@ -11,11 +11,13 @@ REPO_DIR = Path(__file__).parent
 # The repository's main directories
 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 directories if it does not exist yet
 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.py

@@ -14,9 +14,6 @@ from concrete.ml.onnx.convert import get_equivalent_numpy_forward
 from concrete.ml.torch.numpy_module import NumpyModule
 from concrete.ml.version import __version__ as CML_VERSION
 
-# Add a "black and white" filter
-# FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2277
-
 
 class _TorchIdentity(nn.Module):
     """Torch identity model."""
@@ -189,9 +186,6 @@ class Filter:
             self.torch_model = _TorchConv2D(kernel, n_out_channels=3, groups=3)
 
         elif image_filter == "ridge detection":
-            # Make the filter properly grayscaled, as it is commonly used
-            # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/2265
-
             kernel = torch.tensor(
                 [
                     [-1, -1, -1],

filters/black and white/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:63773fc0796c373f650f74b5462515a7bf95c2b591266472f23bf37b0846e935
+oid sha256:9a8fae28225c53cc6e184535a8880187626892461a7f0c25afe322dfaa83f678
 size 388

filters/black and white/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8e3d461f5553416f5f011b25b76afdbb0cd0f57b64a2cdfa34b705049cd796a9
-size 4367
+oid sha256:4253afb0e46be27e5aa2562c14e7bc402fc7aa39cf01a824f1170c5e46ccf9aa
+size 4364

filters/blur/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1b9b07a161dbe6020a0749c2bec2ae9e4215ac25574ac086213a45979a8c229c
+oid sha256:6ead03bc109e6598c90000f580c2d149f19c4ed968dfa302dfea78c131bdad02
 size 391

filters/blur/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:34f388a428c9d9c61d5138ee86c0248a2c495b5e2c361a8839551f20dc0f19f6
-size 7266
+oid sha256:deb92cf90d25f2c4552bb8a2fa0eb0f0cbe0fdc549665f2b63a15a78c0b20d72
+size 7263

filters/identity/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:adbdb74b0fb4dec87aaa065bf9c90940004c5a1f4cb1c0943c0134237b8008cd
+oid sha256:05f39bf162b0ffcd9e4105c371180a621ebdc64fce12a751547424e3a6ec0c0b
 size 378

filters/identity/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5dd04e43447e0ad26bbbc48db67b1aa7fc2993862bdce53eabd2d39793240eec
-size 2562
+oid sha256:9741844b0e166078608310f3812724cf0dd117927457763ecec48e5e510ebc5a
+size 2559

filters/inverted/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:adbdb74b0fb4dec87aaa065bf9c90940004c5a1f4cb1c0943c0134237b8008cd
+oid sha256:05f39bf162b0ffcd9e4105c371180a621ebdc64fce12a751547424e3a6ec0c0b
 size 378

filters/inverted/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:25c281b6d13a16e5e41ee53b929dd33318b24f21ce91e9120aa5ed2065b4e7ab
-size 4182
+oid sha256:e3b413c12bb4d656e739247af1ffb848d57a70a2f57011dd367af66c56cce025
+size 4179

filters/ridge detection/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a5f28e47bbfa1cc51e4212b471ba8fb981af51f7db23a93892ec849a2ec2541f
+oid sha256:b58c440c56ec350092d730eb2e8d2118225eed501fb8dba5df17b96440fe6a08
 size 397

filters/ridge detection/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5a621b96e5a707c67efc53c6ac262676558545c8d173df9cab5138866a46afe3
-size 4483
+oid sha256:91c225472c2d225d3211b54a1f4bca1a7df54ae6543e6d62e7c6a48601dd9e31
+size 4479

filters/rotate/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0f3d46ba45c50191d39587992f3a53827949214481ac05a9673f2481f9038ab4
+oid sha256:05f39bf162b0ffcd9e4105c371180a621ebdc64fce12a751547424e3a6ec0c0b
 size 378

filters/rotate/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5599e84aa1aaf1a908f17dd15c4e221c234dcc6998a17ada305cfc4aaceade6d
-size 4434
+oid sha256:3f4c3cfa73f19125923edca845c70a55eb5a882ddb17fc8f979ed349b845eb6e
+size 4431

filters/sharpen/deployment/client.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e6fe9d09eacf0ffcc5569457fb3e869488a35ce134583c70b60cf5c099c8b46e
+oid sha256:4a13c773025de645be26fbedd7bb9b0926464ad08f992ff089aee90ee58df5f3
 size 396

filters/sharpen/deployment/server.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ab65ba0d811d187405b681fd5210da8c996111a53ffac48de8ff20457ac46949
-size 7314
+oid sha256:65892df870b0e07b717fb971b85c0e6ed2d0acc345db9fec1eb90497b1092374
+size 7311

requirements.txt

@@ -2,4 +2,3 @@ concrete-ml==0.6.1
 gradio==3.11.0
 uvicorn==0.20.0
 fastapi==0.87.0
-jupyter==1.0.0