baseline, wip

app.py

@@ -1,7 +1,437 @@
+"""A local gradio app that detects seizures with EEG using FHE."""
+from PIL import Image
+import os
+import shutil
+import subprocess
+import time
 import gradio as gr
+import numpy
+import requests
+from itertools import chain
 
-def greet(name):
-    return "Hello " + name + "!!"
+from common import (
+    CLIENT_TMP_PATH,
+    SERVER_TMP_PATH,
+    EXAMPLES,
+    INPUT_SHAPE,
+    KEYS_PATH,
+    REPO_DIR,
+    SERVER_URL,
+)
+from client_server_interface import FHEClient
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# 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 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):
+    """Get the client API.
+
+    Args:
+        user_id (int): The current user's ID.
+
+    Returns:
+        FHEClient: The client API.
+    """
+    return FHEClient(
+        key_dir=KEYS_PATH / f"seizure_detection_{user_id}",
+    )
+
+def get_client_file_path(name, user_id):
+    """Get the correct temporary file path for the client.
+
+    Args:
+        name (str): The desired file name.
+        user_id (int): The current user's ID.
+
+    Returns:
+        pathlib.Path: The file path.
+    """
+    return CLIENT_TMP_PATH / f"{name}_seizure_detection_{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():
+    """Generate the private key for seizure detection.
+
+    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)
+
+    # Generate a private key
+    client.generate_private_and_evaluation_keys(force=True)
+
+    # Retrieve the serialized evaluation key
+    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)
+
+    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):
+    """Encrypt the given image for seizure detection.
+
+    Args:
+        user_id (int): The current user's ID.
+        input_image (numpy.ndarray): The image to encrypt.
+
+    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.")
+
+    if input_image.shape[-1] != 3:
+        raise ValueError(f"Input image must have 3 channels (RGB). Current shape: {input_image.shape}")
+
+    # Resize the image if it hasn't the shape (224, 224, 3)
+    if input_image.shape != (224, 224, 3):
+        input_image_pil = Image.fromarray(input_image)
+        input_image_pil = input_image_pil.resize((224, 224))
+        input_image = numpy.array(input_image_pil)
+
+    # Retrieve the client API
+    client = get_client(user_id)
+
+    # 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)
+
+    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 (resize_img(input_image), encrypted_image_short)
+
+def send_input(user_id):
+    """Send the encrypted input image as well as the evaluation key to the server.
+
+    Args:
+        user_id (int): The current user's ID.
+    """
+    # Get the evaluation key path
+    evaluation_key_path = get_client_file_path("evaluation_key", user_id)
+
+    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)
+
+    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,
+    }
+
+    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):
+    """Apply the seizure detection model on the encrypted image previously sent using FHE.
+
+    Args:
+        user_id (int): The current user's ID.
+    """
+    data = {
+        "user_id": user_id,
+    }
+
+    # 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):
+    """Retrieve the encrypted output (boolean).
+
+    Args:
+        user_id (int): The current user's ID.
+
+    Returns:
+        encrypted_output_short (bytes): A representation of the encrypted result.
+
+    """
+    data = {
+        "user_id": user_id,
+    }
+
+    # Retrieve the encrypted output
+    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)
+
+            with encrypted_output_path.open("wb") as encrypted_output_file:
+                encrypted_output_file.write(encrypted_output)
+
+            # Create a truncated version of the encrypted output for display
+            encrypted_output_short = shorten_bytes_object(encrypted_output)
+
+            return encrypted_output_short
+        else:
+            raise gr.Error("Please wait for the FHE execution to be completed.")
+
+def decrypt_output(user_id):
+    """Decrypt the result.
+
+    Args:
+        user_id (int): The current user's ID.
+
+    Returns:
+        bool: The decrypted output (True if seizure detected, False otherwise)
+
+    """
+    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)
+
+    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 = encrypted_output_file.read()
+
+    # Retrieve the client API
+    client = get_client(user_id)
+
+    # Deserialize, decrypt and post-process the encrypted output
+    decrypted_output = client.deserialize_decrypt_post_process(encrypted_output)
+
+    return "Seizure detected" if decrypted_output else "No seizure detected"
+
+def resize_img(img, width=256, height=256):
+    """Resize the image."""
+    if img.dtype != numpy.uint8:
+        img = img.astype(numpy.uint8)
+    img_pil = Image.fromarray(img)
+    # Resize the image
+    resized_img_pil = img_pil.resize((width, height))
+    # Convert back to a NumPy array
+    return numpy.array(resized_img_pil)
+
+demo = gr.Blocks()
+
+print("Starting the demo...")
+with demo:
+    gr.Markdown(
+        """
+        <h1 align="center">Seizure Detection on Encrypted EEG Data Using Fully Homomorphic Encryption</h1>
+        """
+    )
+
+    gr.Markdown("## Client side")
+    gr.Markdown("### Step 1: Upload an EEG image. ")
+    gr.Markdown(
+        f"The image will automatically be resized to shape (224x224). "
+        "The image here, however, is displayed in its original resolution."
+    )
+    with gr.Row():
+        input_image = gr.Image(
+            value=None, label="Upload an EEG image here.", height=256,
+            width=256, sources="upload", interactive=True,
+        )
+
+        examples = gr.Examples(
+            examples=EXAMPLES, inputs=[input_image], examples_per_page=5, label="Examples to use."
+        )
+
+    gr.Markdown("### Step 2: 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 3: 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 seizure "
+        "detection directly over encrypted values. Once the computation is finished, the server returns "
+        "the encrypted results to the client."
+    )
+
+    gr.Markdown("### Step 4: 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 5: 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 6: Receive the encrypted output from the server.")
+    get_output_button = gr.Button("Receive the encrypted output from the server.")
+
+    with gr.Row():
+        encrypted_output = gr.Textbox(
+            label="Encrypted output representation:",
+            max_lines=2,
+            interactive=False
+        )
+
+    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 the detection result."
+    )
+
+    gr.Markdown("### Step 7: Decrypt the output.")
+    decrypt_button = gr.Button("Decrypt the output")
+
+    with gr.Row():
+        decrypted_output = gr.Textbox(
+            label="Seizure detection result:",
+            interactive=False
+        )
+
+    # Button to generate the private key
+    keygen_button.click(
+        keygen,
+        outputs=[user_id, keygen_checkbox],
+    )
+
+    # Button to encrypt inputs on the client side
+    encrypt_button.click(
+        encrypt,
+        inputs=[user_id, input_image],
+        outputs=[input_image, encrypted_input],
+    )
+
+    # Button to send the encodings to the server using post method
+    send_input_button.click(
+        send_input, inputs=[user_id], outputs=[send_input_checkbox]
+    )
+
+    # Button to send the encodings to the server using post method
+    execute_fhe_button.click(run_fhe, inputs=[user_id], outputs=[fhe_execution_time])
+
+    # Button to send the encodings to the server using post method
+    get_output_button.click(
+        get_output,
+        inputs=[user_id],
+        outputs=[encrypted_output]
+    )
+
+    # Button to decrypt the output on the client side
+    decrypt_button.click(
+        decrypt_output,
+        inputs=[user_id],
+        outputs=[decrypted_output],
+    )
+
+    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)

client_server_interface.py

@@ -0,0 +1,150 @@
+"Client-server interface custom implementation for seizure detection models."
+
+from concrete import fhe
+
+from seizure_detection import SeizureDetector
+
+
+class FHEServer:
+    """Server interface to run a FHE circuit for seizure detection."""
+
+    def __init__(self, model_path):
+        """Initialize the FHE interface.
+
+        Args:
+            model_path (Path): The path to the directory where the circuit is saved.
+        """
+        self.model_path = model_path
+
+        # Load the FHE circuit
+        self.server = fhe.Server.load(self.model_path / "server.zip")
+
+    def run(self, serialized_encrypted_image, serialized_evaluation_keys):
+        """Run seizure detection on the server over an encrypted image.
+
+        Args:
+            serialized_encrypted_image (bytes): The encrypted and serialized image.
+            serialized_evaluation_keys (bytes): The serialized evaluation keys.
+
+        Returns:
+            bytes: The encrypted boolean output indicating seizure detection.
+        """
+        # Deserialize the encrypted input image and the evaluation keys
+        encrypted_image = fhe.Value.deserialize(serialized_encrypted_image)
+        evaluation_keys = fhe.EvaluationKeys.deserialize(serialized_evaluation_keys)
+
+        # Execute the seizure detection in FHE
+        encrypted_output = self.server.run(encrypted_image, evaluation_keys=evaluation_keys)
+
+        # Serialize the encrypted output
+        serialized_encrypted_output = encrypted_output.serialize()
+
+        return serialized_encrypted_output
+
+
+class FHEDev:
+    """Development interface to save and load the seizure detection model."""
+
+    def __init__(self, seizure_detector, model_path):
+        """Initialize the FHE interface.
+
+        Args:
+            seizure_detector (SeizureDetector): The seizure detection model to use in the FHE interface.
+            model_path (str): The path to the directory where the circuit is saved.
+        """
+
+        self.seizure_detector = seizure_detector
+        self.model_path = model_path
+
+        self.model_path.mkdir(parents=True, exist_ok=True)
+
+    def save(self):
+        """Export all needed artifacts for the client and server interfaces."""
+
+        assert self.seizure_detector.fhe_circuit is not None, (
+            "The model must be compiled before saving it."
+        )
+
+        # Save the circuit for the server, using the via_mlir in order to handle cross-platform
+        # execution
+        path_circuit_server = self.model_path / "server.zip"
+        self.seizure_detector.fhe_circuit.server.save(path_circuit_server, via_mlir=True)
+
+        # Save the circuit for the client
+        path_circuit_client = self.model_path / "client.zip"
+        self.seizure_detector.fhe_circuit.client.save(path_circuit_client)
+
+
+class FHEClient:
+    """Client interface to encrypt and decrypt FHE data associated to a SeizureDetector."""
+
+    def __init__(self, model_path, key_dir=None):
+        """Initialize the FHE interface.
+
+        Args:
+            model_path (Path): The path to the directory where the circuit is saved.
+            key_dir (Path): The path to the directory where the keys are stored. Default to None.
+        """
+        self.model_path = model_path
+        self.key_dir = key_dir
+
+        # If model_path does not exist raise
+        assert model_path.exists(), f"{model_path} does not exist. Please specify a valid path."
+
+        # Load the client
+        self.client = fhe.Client.load(self.model_path / "client.zip", self.key_dir)
+
+        # Instantiate the seizure detector
+        self.seizure_detector = SeizureDetector()
+
+    def generate_private_and_evaluation_keys(self, force=False):
+        """Generate the private and evaluation keys.
+
+        Args:
+            force (bool): If True, regenerate the keys even if they already exist.
+        """
+        self.client.keygen(force)
+
+    def get_serialized_evaluation_keys(self):
+        """Get the serialized evaluation keys.
+
+        Returns:
+            bytes: The evaluation keys.
+        """
+        return self.client.evaluation_keys.serialize()
+
+    def encrypt_serialize(self, input_image):
+        """Encrypt and serialize the input image in the clear.
+
+        Args:
+            input_image (numpy.ndarray): The image to encrypt and serialize.
+
+        Returns:
+            bytes: The pre-processed, encrypted and serialized image.
+        """
+        # Encrypt the image
+        encrypted_image = self.client.encrypt(input_image)
+
+        # Serialize the encrypted image to be sent to the server
+        serialized_encrypted_image = encrypted_image.serialize()
+        return serialized_encrypted_image
+
+    def deserialize_decrypt_post_process(self, serialized_encrypted_output):
+        """Deserialize, decrypt and post-process the output in the clear.
+
+        Args:
+            serialized_encrypted_output (bytes): The serialized and encrypted output.
+
+        Returns:
+            bool: The decrypted and deserialized boolean indicating seizure detection.
+        """
+        # Deserialize the encrypted output
+        encrypted_output = fhe.Value.deserialize(serialized_encrypted_output)
+
+        # Decrypt the output
+        output = self.client.decrypt(encrypted_output)
+
+        # Post-process the output (if needed)
+        seizure_detected = self.seizure_detector.post_processing(output)
+
+        return seizure_detected

common.py

@@ -0,0 +1,29 @@
+"All the constants used in this repo."
+
+from pathlib import Path
+
+# This repository's directory
+REPO_DIR = Path(__file__).parent
+
+# This repository's main necessary folders
+FILTERS_PATH = REPO_DIR / "filters"
+KEYS_PATH = REPO_DIR / ".fhe_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)
+CLIENT_TMP_PATH.mkdir(exist_ok=True)
+SERVER_TMP_PATH.mkdir(exist_ok=True)
+
+# The input images' shape. Images with different input shapes will be cropped and resized by Gradio
+INPUT_SHAPE = (224, 224)
+
+# Retrieve the input examples directory
+INPUT_EXAMPLES_DIR = REPO_DIR / "input_examples"
+
+# List of all image examples suggested in the demo
+EXAMPLES = [str(image) for image in INPUT_EXAMPLES_DIR.glob("**/*")]
+
+# Store the server's URL
+SERVER_URL = "http://localhost:8000/"

requirements.txt

@@ -0,0 +1,2 @@
+concrete-ml==1.1.0
+gradio

server.py

@@ -0,0 +1,97 @@
+"""Server that will listen for GET and POST requests from the client."""
+
+import time
+from typing import List
+from fastapi import FastAPI, File, Form, UploadFile
+from fastapi.responses import JSONResponse, Response
+
+from common import SERVER_TMP_PATH
+from client_server_interface import FHEServer
+
+# Load the server object for seizure detection
+FHE_SERVER = FHEServer(model_path="path/to/seizure_detection_model")
+
+def get_server_file_path(name, user_id):
+    """Get the correct temporary file path for the server.
+
+    Args:
+        name (str): The desired file name.
+        user_id (int): The current user's ID.
+
+    Returns:
+        pathlib.Path: The file path.
+    """
+    return SERVER_TMP_PATH / f"{name}_seizure_detection_{user_id}"
+
+
+# Initialize an instance of FastAPI
+app = FastAPI()
+
+# Define the default route
+@app.get("/")
+def root():
+    return {"message": "Welcome to Your Seizure Detection FHE Server!"}
+
+
+@app.post("/send_input")
+def send_input(
+    user_id: str = Form(),
+    files: List[UploadFile] = File(),
+):
+    """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", user_id)
+    evaluation_key_path = get_server_file_path("evaluation_key", user_id)
+
+    # Write the files using the above paths
+    with encrypted_image_path.open("wb") as encrypted_image, evaluation_key_path.open(
+        "wb"
+    ) as evaluation_key:
+        encrypted_image.write(files[0].file.read())
+        evaluation_key.write(files[1].file.read())
+
+
+@app.post("/run_fhe")
+def run_fhe(
+    user_id: str = Form(),
+):
+    """Execute seizure detection 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", user_id)
+    evaluation_key_path = get_server_file_path("evaluation_key", user_id)
+
+    # Read the files using the above paths
+    with encrypted_image_path.open("rb") as encrypted_image_file, evaluation_key_path.open(
+        "rb"
+    ) as evaluation_key_file:
+        encrypted_image = encrypted_image_file.read()
+        evaluation_key = evaluation_key_file.read()
+
+    # Run the FHE execution
+    start = time.time()
+    encrypted_output = FHE_SERVER.run(encrypted_image, evaluation_key)
+    fhe_execution_time = round(time.time() - start, 2)
+
+    # Retrieve the encrypted output path
+    encrypted_output_path = get_server_file_path("encrypted_output", user_id)
+
+    # Write the file using the above path
+    with encrypted_output_path.open("wb") as encrypted_output_file:
+        encrypted_output_file.write(encrypted_output)
+
+    return JSONResponse(content=fhe_execution_time)
+
+
+@app.post("/get_output")
+def get_output(
+    user_id: str = Form(),
+):
+    """Retrieve the encrypted output."""
+    # Retrieve the encrypted output path
+    encrypted_output_path = get_server_file_path("encrypted_output", user_id)
+
+    # Read the file using the above path
+    with encrypted_output_path.open("rb") as encrypted_output_file:
+        encrypted_output = encrypted_output_file.read()
+
+    return Response(encrypted_output)