File size: 18,796 Bytes
21c7197
 
 
 
 
 
 
 
 
5369c66
 
21c7197
 
 
5369c66
21c7197
 
 
 
 
 
 
7b32412
21c7197
 
 
 
 
 
4bee342
3cf0931
 
 
4bee342
d07e70b
3cf0931
5369c66
d07e70b
 
7b32412
 
d07e70b
5369c66
 
 
 
 
 
 
 
 
 
 
d07e70b
 
 
21c7197
 
 
 
 
 
 
 
 
 
 
7b32412
21c7197
 
 
 
 
 
 
4bee342
21c7197
 
 
 
4bee342
21c7197
 
7b32412
21c7197
7b32412
2a040cc
 
5369c66
21c7197
 
 
4bee342
21c7197
 
 
 
 
4bee342
21c7197
 
 
 
4bee342
21c7197
 
 
 
 
5369c66
 
21c7197
 
5369c66
21c7197
 
5369c66
 
21c7197
 
 
5369c66
 
 
 
 
21c7197
 
5369c66
 
21c7197
5369c66
 
21c7197
5369c66
21c7197
 
 
4bee342
21c7197
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
 
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
 
 
4bee342
21c7197
 
 
 
2eb250a
21c7197
 
4bee342
21c7197
 
 
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
d07e70b
 
 
21c7197
4bee342
21c7197
7b32412
1974e22
21c7197
 
 
4bee342
21c7197
 
 
 
 
 
 
2eb250a
21c7197
 
4bee342
21c7197
 
 
 
4bee342
21c7197
 
4bee342
21c7197
 
 
 
4bee342
21c7197
 
 
 
 
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4bee342
21c7197
 
 
 
4bee342
21c7197
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
 
 
 
 
 
4bee342
21c7197
 
 
 
4bee342
21c7197
 
 
 
 
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
d07e70b
 
21c7197
 
4bee342
21c7197
 
d07e70b
21c7197
d07e70b
4bee342
21c7197
2eb250a
21c7197
 
 
 
4bee342
21c7197
 
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
4bee342
21c7197
 
7b32412
21c7197
 
 
 
 
 
 
 
 
 
 
 
f9f9bfd
21c7197
75bc44e
21c7197
f9f9bfd
 
 
 
 
 
 
 
 
b3b800b
502bbf4
 
b3b800b
 
 
21c7197
 
 
 
 
502bbf4
21c7197
502bbf4
 
 
21c7197
 
 
2eb250a
21c7197
 
 
 
 
 
502bbf4
4bee342
21c7197
 
 
d07e70b
21c7197
 
502bbf4
 
21c7197
 
d07e70b
502bbf4
b08dbdd
21c7197
b08dbdd
21c7197
 
d07e70b
 
502bbf4
b08dbdd
d07e70b
b08dbdd
d07e70b
 
 
21c7197
 
 
 
 
 
 
 
502bbf4
b08dbdd
 
21c7197
502bbf4
b08dbdd
 
 
 
21c7197
502bbf4
d07e70b
502bbf4
d07e70b
 
 
b08dbdd
21c7197
b08dbdd
d07e70b
14c30ad
 
 
 
d07e70b
 
21c7197
 
502bbf4
21c7197
 
 
502bbf4
d07e70b
 
 
 
 
21c7197
 
 
 
 
 
 
14c30ad
 
21c7197
 
 
14c30ad
 
 
 
21c7197
 
 
 
 
4bee342
2eb250a
21c7197
 
 
 
 
4bee342
2eb250a
21c7197
 
 
 
4bee342
21c7197
 
 
4bee342
21c7197
 
 
d07e70b
4bee342
2eb250a
21c7197
 
 
 
 
4bee342
21c7197
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
"""A local gradio app that filters images using FHE."""

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,
    REPO_DIR,
    SERVER_URL,
)
from client_server_interface import FHEClient

# 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, filter_name):
    """Decrypt the encrypted output using a different private key.
    """
    # Retrieve the filter's deployment path
    filter_path = FILTERS_PATH / f"{filter_name}/deployment"

    # Instantiate the client interface and generate a new private key
    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


def shorten_bytes_object(bytes_object, limit=500):
    """Shorten the input bytes object to a given length.

    Encrypted data is too large for displaying it in the browser using Gradio. This function
    provides a shorten representation of it.

    Args:
        bytes_object (bytes): The input to shorten
        limit (int): The length to consider. Default to 500.

    Returns:
        str: Hexadecimal string shorten representation of the input byte object. 

    """
    # Define a shift for better display
    shift = 100
    return bytes_object[shift : limit + shift].hex()


def get_client(user_id, filter_name):
    """Get the client API.

    Args:
        user_id (int): The current user's ID.
        filter_name (str): The filter chosen by the user

    Returns:
        FHEClient: The client API.
    """
    return FHEClient(
        FILTERS_PATH / f"{filter_name}/deployment", 
        filter_name,
        key_dir=KEYS_PATH / f"{filter_name}_{user_id}", 
    )


def get_client_file_path(name, user_id, filter_name):
    """Get the correct temporary file path for the client.

    Args:
        name (str): The desired file name.
        user_id (int): The current user's ID.
        filter_name (str): The filter chosen by the user

    Returns:
        pathlib.Path: The file path.
    """
    return CLIENT_TMP_PATH / f"{name}_{filter_name}_{user_id}"


def clean_temporary_files(n_keys=20):
    """Clean keys and encrypted images.

    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 and associated files to be stored. Default to 20.

    """
    # 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(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
    client_files = CLIENT_TMP_PATH.iterdir()
    server_files = SERVER_TMP_PATH.iterdir()

    # 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 user_id in file.name:
                file.unlink()


def keygen(filter_name):
    """Generate the private key associated to a filter.

    Args:
        filter_name (str): The current filter to consider.

    Returns:
        (user_id, True) (Tuple[int, bool]): The current user's ID and a boolean used for visual display.

    """
    # Clean temporary files
    clean_temporary_files()

    # Create an ID for the current user
    user_id = numpy.random.randint(0, 2**32)

    # Retrieve the client API
    client = get_client(user_id, filter_name)

    # Generate a private key
    client.generate_private_and_evaluation_keys(force=True)

    # Retrieve the serialized evaluation key. In this case, as circuits are fully leveled, this
    # evaluation key is empty. However, for software reasons, it is still needed for proper FHE
    # execution
    evaluation_key = client.get_serialized_evaluation_keys()

    # 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, filter_name)

    with evaluation_key_path.open("wb") as evaluation_key_file:
        evaluation_key_file.write(evaluation_key)

    return (user_id, True)


def encrypt(user_id, input_image, filter_name):
    """Encrypt the given image for a specific user and filter.

    Args:
        user_id (int): The current user's ID.
        input_image (numpy.ndarray): The image to encrypt.
        filter_name (str): The current filter to consider.

    Returns:
        (input_image, encrypted_image_short) (Tuple[bytes]): The encrypted image and one of its
        representation.

    """
    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, filter_name)

    # Pre-process, encrypt and serialize the image
    encrypted_image = client.encrypt_serialize(input_image)

    # 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, filter_name)

    with encrypted_image_path.open("wb") as encrypted_image_file:
        encrypted_image_file.write(encrypted_image)

    # Create a truncated version of the encrypted image for display
    encrypted_image_short = shorten_bytes_object(encrypted_image)

    return (input_image, encrypted_image_short)


def send_input(user_id, filter_name):
    """Send the encrypted input image as well as the evaluation key to the server.

    Args:
        user_id (int): The current user's ID.
        filter_name (str): The current filter to consider.
    """
    # Get the evaluation key path
    evaluation_key_path = get_client_file_path("evaluation_key", user_id, filter_name)

    if user_id == "" or not evaluation_key_path.is_file():
        raise gr.Error("Please generate the private key first.")

    encrypted_input_path = get_client_file_path("encrypted_image", user_id, filter_name)

    if not encrypted_input_path.is_file():
        raise gr.Error("Please generate the private key and then encrypt an image first.")

    # Define the data and files to post
    data = {
        "user_id": user_id,
        "filter": filter_name,
    }

    files = [
        ("files", open(encrypted_input_path, "rb")),
        ("files", open(evaluation_key_path, "rb")),
    ]

    # Send the encrypted input image and evaluation key to the server
    url = SERVER_URL + "send_input"
    with requests.post(
        url=url,
        data=data,
        files=files,
    ) as response:
        return response.ok


def run_fhe(user_id, filter_name):
    """Apply the filter on the encrypted image previously sent using FHE.

    Args:
        user_id (int): The current user's ID.
        filter_name (str): The current filter to consider.
    """
    data = {
        "user_id": user_id,
        "filter": filter_name,
    }

    # Trigger the FHE execution on the encrypted image previously sent
    url = SERVER_URL + "run_fhe"
    with requests.post(
        url=url,
        data=data,
    ) as response:
        if response.ok:
            return response.json()
        else:
            raise gr.Error("Please wait for the input image to be sent to the server.")


def get_output(user_id, filter_name):
    """Retrieve the encrypted output image.

    Args:
        user_id (int): The current user's ID.
        filter_name (str): The current filter to consider.

    Returns:
        encrypted_output_image_short (bytes): A representation of the encrypted result.

    """
    data = {
        "user_id": user_id,
        "filter": filter_name,
    }

    # Retrieve the encrypted output image
    url = SERVER_URL + "get_output"
    with requests.post(
        url=url,
        data=data,
    ) as response:
        if response.ok:
            encrypted_output = response.content

            # 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, filter_name)

            with encrypted_output_path.open("wb") as encrypted_output_file:
                encrypted_output_file.write(encrypted_output)

            # Decrypt the image using a different (wrong) key for display
            output_image_representation = decrypt_output_with_wrong_key(encrypted_output, filter_name)

            return output_image_representation
        else:
            raise gr.Error("Please wait for the FHE execution to be completed.")


def decrypt_output(user_id, filter_name):
    """Decrypt the result.

    Args:
        user_id (int): The current user's ID.
        filter_name (str): The current filter to consider.

    Returns:
        (output_image, False, False) ((Tuple[numpy.ndarray, bool, bool]): The decrypted output, as
            well as two booleans used for resetting Gradio checkboxes

    """
    if user_id == "":
        raise gr.Error("Please generate the private key first.")

    # Get the encrypted output path
    encrypted_output_path = get_client_file_path("encrypted_output", user_id, filter_name)

    if not encrypted_output_path.is_file():
        raise gr.Error("Please run the FHE execution first.")

    # Load the encrypted output as bytes
    with encrypted_output_path.open("rb") as encrypted_output_file:
        encrypted_output_image = encrypted_output_file.read()

    # Retrieve the client API
    client = get_client(user_id, filter_name)

    # Deserialize, decrypt and post-process the encrypted output
    output_image = client.deserialize_decrypt_post_process(encrypted_output_image)

    return output_image, False, False


demo = gr.Blocks()


print("Starting the demo...")
with demo:
    gr.Markdown(
        """
        <p align="center">
            <img width=200 src="https://user-images.githubusercontent.com/5758427/197816413-d9cddad3-ba38-4793-847d-120975e1da11.png">
        </p>
        <h1 align="center">Image Filtering On Encrypted Data Using Fully Homomorphic Encryption</h1>
        <p align="center">
            <a href="https://github.com/zama-ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197972109-faaaff3e-10e2-4ab6-80f5-7531f7cfb08f.png">Concrete-ML</a>

            <a href="https://docs.zama.ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197976802-fddd34c5-f59a-48d0-9bff-7ad1b00cb1fb.png">Documentation</a>

            <a href="https://zama.ai/community"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197977153-8c9c01a7-451a-4993-8e10-5a6ed5343d02.png">Community</a>

            <a href="https://twitter.com/zama_fhe"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197975044-bab9d199-e120-433b-b3be-abd73b211a54.png">@zama_fhe</a>
        </p>
        <p align="center">
            <img src="https://user-images.githubusercontent.com/56846628/219605302-5baafac4-cf6f-4f06-9a96-91cef2b84a63.png" width="70%" height="70%">
        </p>
        <p align="center">
            Test the app below, review
            <a href="https://zama.ai/post/encrypted-image-filtering-using-homomorphic-encryption"> our tutorial</a>
            , and try the build for yourself!
        </p>
        """
    )

    gr.Markdown("## Client side")
    gr.Markdown("### Step 1: Upload an image. ")
    gr.Markdown(
        f"The image will automatically be resized to shape ({INPUT_SHAPE[0]}x{INPUT_SHAPE[1]}). "
        "The image here, however, is displayed in its original resolution. The true image used "
        "in this demo can be seen in Step 8."
    )
    with gr.Row():
        input_image = gr.Image(
            value=None, label="Upload an image here.", shape=INPUT_SHAPE, source="upload", interactive=True,
        )

        examples = gr.Examples(
            examples=EXAMPLES, inputs=[input_image], examples_per_page=5, label="Examples to use."
        )

    gr.Markdown("### Step 2: Choose your filter.")
    filter_name = gr.Dropdown(
        choices=AVAILABLE_FILTERS, value="inverted", label="Choose your filter", interactive=True
    )

    gr.Markdown("#### Notes")
    gr.Markdown(
        """
        - The private key is used to encrypt and decrypt the data and will never be shared.
        - No public key is required for these filter operators.
        """
    )
    
    gr.Markdown("### Step 3: Generate the private key.")
    keygen_button = gr.Button("Generate the private key.")

    with gr.Row():
        keygen_checkbox = gr.Checkbox(label="Private key generated:", interactive=False)

    user_id = gr.Textbox(label="", max_lines=2, interactive=False, visible=False)
    
    gr.Markdown("### Step 4: Encrypt the image using FHE.")
    encrypt_button = gr.Button("Encrypt the image using FHE.")

    with gr.Row():
        encrypted_input = gr.Textbox(
            label="Encrypted input representation:", max_lines=2, interactive=False
        )

    gr.Markdown("## Server side")
    gr.Markdown(
        "The encrypted value is received by the server. The server can then compute the filter "
        "directly over encrypted values. Once the computation is finished, the server returns "
        "the encrypted results to the client."
    )

    gr.Markdown("### 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.")
    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
    )

    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."
    )
    get_output_button = gr.Button("Receive the encrypted output image from the server.")

    with gr.Row():
        encrypted_output_representation = gr.Image(
            label=f"Encrypted output representation ({INPUT_SHAPE[0]}x{INPUT_SHAPE[1]}):", 
            interactive=False,
            height=256,
            width=256,
        )

    gr.Markdown("## Client side")
    gr.Markdown(
        "The encrypted output is sent back to the client, who can finally decrypt it with the "
        "private key. Only the client is aware of the original image and its transformed version."
    )

    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():
        original_image = gr.Image(
            input_image.value,
            label=f"Input image ({INPUT_SHAPE[0]}x{INPUT_SHAPE[1]}):",
            interactive=False,
            height=256,
            width=256,   
        )

        output_image = gr.Image(
            label=f"Output image ({INPUT_SHAPE[0]}x{INPUT_SHAPE[1]}):", 
            interactive=False,
            height=256,
            width=256,
        )

    # Button to generate the private key
    keygen_button.click(
        keygen,
        inputs=[filter_name],
        outputs=[user_id, keygen_checkbox],
    )

    # Button to encrypt inputs on the client side
    encrypt_button.click(
        encrypt,
        inputs=[user_id, input_image, filter_name],
        outputs=[original_image, encrypted_input],
    )

    # Button to send the encodings to the server using post method
    send_input_button.click(
        send_input, inputs=[user_id, filter_name], outputs=[send_input_checkbox]
    )

    # Button to send the encodings to the server using post method
    execute_fhe_button.click(run_fhe, inputs=[user_id, filter_name], outputs=[fhe_execution_time])

    # Button to send the encodings to the server using post method
    get_output_button.click(
        get_output, 
        inputs=[user_id, filter_name], 
        outputs=[encrypted_output_representation]
    )

    # Button to decrypt the output on the client side
    decrypt_button.click(
        decrypt_output,
        inputs=[user_id, filter_name],
        outputs=[output_image, keygen_checkbox, send_input_checkbox],
    )

    gr.Markdown(
        "The app was built with [Concrete-ML](https://github.com/zama-ai/concrete-ml), a "
        "Privacy-Preserving Machine Learning (PPML) open-source set of tools by [Zama](https://zama.ai/). "
        "Try it yourself and don't forget to star on Github &#11088;."
    )

demo.launch(share=False)