Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
 import os
+import numpy as np
 from PIL import Image
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 from cryptography.hazmat.backends import default_backend
+import cv2
 
 class CBCEncryption:
     def __init__(self, key, iv):
@@ -22,9 +24,22 @@ class CBCEncryption:
     def finalize_decrypt(self):
         return self.decryptor.finalize()
 
+def add_gaussian_noise(image, mean=0, var=0.01):
+    noise = np.random.normal(mean, var**0.5, image.shape)
+    noisy_image = image + noise
+    return noisy_image
+
+def remove_gaussian_noise(image, mean=0, var=0.01):
+    # This is a simplified approach to remove noise.
+    # In practice, you might need a more sophisticated method to estimate the noise statistics.
+    return image - np.random.normal(mean, var**0.5, image.shape)
+
 def EncryptImage(encryption, image, key):
     image = Image.open(image)
-    image.save('temp.bmp')
+    image_np = np.array(image)
+    noisy_image_np = add_gaussian_noise(image_np)
+    noisy_image = Image.fromarray(noisy_image_np.astype(np.uint8))
+    noisy_image.save('temp.bmp')
     with open('temp.bmp', 'rb') as reader:
         with open('encrypted.bmp', 'wb') as writer:
             image_data = reader.read()
@@ -42,7 +57,10 @@ def DecryptImage(decryption, image, key):
             header, body = image_data[:54], image_data[54:]
             body = decryption.decrypt(body)
             unpadded_body = unpad_data(body)
-            writer.write(header + unpadded_body)
+            decrypted_image_np = np.frombuffer(unpadded_body, dtype=np.uint8).reshape(image.shape)
+            original_image_np = remove_gaussian_noise(decrypted_image_np)
+            original_image = Image.fromarray(original_image_np.astype(np.uint8))
+            original_image.save('decrypted.bmp')
     return 'decrypted.bmp'
 
 def unpad_data(data):
@@ -78,7 +96,7 @@ def main():
         action, # Single function to wrap
         ["image", "text", "text"], # Input types: image, key, action_type
         "image", # Output type
-        title="AES CBC Encryption/Decryption",
+        title="AES CBC Encryption/Decryption with Gaussian Noise",
         description="Upload an image, a key, and choose Encrypt or Decrypt to process the image."
     )