decrypt

app.py

@@ -13,6 +13,10 @@ import requests
 from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry
 import logging
+import numpy as np
+import seaborn as sns
+import io
+import matplotlib.pyplot as plt
 
 from common import (
     CLIENT_TMP_PATH,
@@ -310,36 +314,58 @@ def get_output(user_id):
         else:
             raise gr.Error("Please wait for the FHE execution to be completed.")
 
-def decrypt_output(user_id):
-    """Decrypt the result.
+def decrypt_output(user_id, encrypted_output):
+    """Decrypt the output of the seizure detection."""
+    if user_id == "":
+        raise gr.Error("Please generate the private key first.")
 
-    Args:
-        user_id (int): The current user's ID.
+    if encrypted_output is None:
+        raise gr.Error("Please run the FHE computation first.")
 
-    Returns:
-        bool: The decrypted output (True if seizure detected, False otherwise)
+    # Retrieve the client API
+    client = get_client(user_id)
 
-    """
-    if user_id == "":
-        raise gr.Error("Please generate the private key first.")
+    try:
+        # Deserialize and decrypt the output
+        decrypted_output = client.deserialize_decrypt_post_process(encrypted_output)
 
-    # Get the encrypted output path
-    encrypted_output_path = get_client_file_path("encrypted_output", user_id)
+        # Reshape the output to match the expected shape
+        decrypted_output = decrypted_output.reshape(1, 1, 32, 32, 3)
 
-    if not encrypted_output_path.is_file():
-        raise gr.Error("Please run the FHE execution first.")
+        # Convert the output to a probability
+        probability = sigmoid(decrypted_output[0, 0, 0, 0, 0])
 
-    # Load the encrypted output as bytes
-    with encrypted_output_path.open("rb") as encrypted_output_file:
-        encrypted_output = encrypted_output_file.read()
+        # Create a heatmap
+        heatmap = create_heatmap(decrypted_output[0, 0, :, :, 0])
 
-    # Retrieve the client API
-    client = get_client(user_id)
+        return probability, heatmap
 
-    # Deserialize, decrypt and post-process the encrypted output
-    decrypted_output = client.deserialize_decrypt_post_process(encrypted_output)
+    except Exception as e:
+        logger.error(f"Error in decrypt_output: {str(e)}")
+        raise gr.Error(f"Decryption failed: {str(e)}")
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
 
-    return "Seizure detected" if decrypted_output else "No seizure detected"
+def create_heatmap(data):
+    # Normalize the data
+    data_normalized = (data - np.min(data)) / (np.max(data) - np.min(data))
+    
+    # Create a heatmap
+    plt.figure(figsize=(6, 6))
+    sns.heatmap(data_normalized, cmap='YlOrRd', cbar=True)
+    plt.title('Activation Heatmap')
+    plt.axis('off')
+    
+    # Save the plot to a BytesIO object
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    
+    # Clear the current figure
+    plt.clf()
+    
+    return buf
 
 def resize_img(img, width=256, height=256):
     """Resize the image."""
@@ -430,10 +456,15 @@ with demo:
     decrypt_button = gr.Button("Decrypt the output")
 
     with gr.Row():
-        decrypted_output = gr.Textbox(
-            label="Seizure detection result:",
-            interactive=False
-        )
+        probability_output = gr.Number(label="Seizure Probability")
+        heatmap_output = gr.Image(label="Activation Heatmap")
+
+    # Update the decrypt button click event
+    decrypt_button.click(
+        decrypt_output,
+        inputs=[user_id, encrypted_output],
+        outputs=[probability_output, heatmap_output]
+    )
 
     # Button to generate the private key
     keygen_button.click(