Update app.py

app.py

@@ -1,106 +1,75 @@
-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):
-        self.cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
-        self.encryptor = self.cipher.encryptor()
-        self.decryptor = self.cipher.decryptor()
-
-    def encrypt(self, image):
-        return self.encryptor.update(image)
-
-    def decrypt(self, image):
-        return self.decryptor.update(image)
-
-    def finalize_encrypt(self):
-        return self.encryptor.finalize()
-
-    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_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()
-            header, body = image_data[:54], image_data[54:]
-            body += b'\x35' * (16 - (len(body) % 16))
-            body = encryption.encrypt(body) + encryption.finalize_encrypt()
-            writer.write(header + body)
-    os.remove('temp.bmp')
-    return 'encrypted.bmp'
-
-def DecryptImage(decryption, image, key):
-    with open(image, 'rb') as reader:
-        with open('decrypted.bmp', 'wb') as writer:
-            image_data = reader.read()
-            header, body = image_data[:54], image_data[54:]
-            body = decryption.decrypt(body)
-            unpadded_body = unpad_data(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):
-    padding_byte = data[-1]
-    padding_length = padding_byte if padding_byte < 16 else 16
-    unpadded_data = data[:-padding_length]
-    return unpadded_data
-
-def main():
-    def action(image, key, action_type):
-        if len(key) > 32:
-            return 'Key is too long. Maximum key length is 32 characters.'
-        key = key.encode('utf-8')
-        key = key.ljust(32, b'\x35')
-
-        iv = key[:16]
-        iv = bytearray(iv)
-        for i in range(16):
-            iv[i] = iv[i] ^ 0x35
-        iv = bytes(iv)
-
-        AesCbc = CBCEncryption(key, iv)
-        if action_type == 'Encrypt':
-            output_path = EncryptImage(AesCbc, image, key)
-        elif action_type == 'Decrypt':
-            output_path = DecryptImage(AesCbc, image, key)
-        else:
-            return 'Invalid action type. Choose either Encrypt or Decrypt.'
-        return output_path
-
-    # Define the Gradio interface
-    iface = gr.Interface(
-        action, # Single function to wrap
-        ["image", "text", "text"], # Input types: image, key, action_type
-        "image", # Output type
-        title="AES CBC Encryption/Decryption with Gaussian Noise",
-        description="Upload an image, a key, and choose Encrypt or Decrypt to process the image."
-    )
-
-    iface.launch()
-
-if __name__ == '__main__':
-    main()
+import numpy as np
+import gradio as gr
+from cryptography.fernet import Fernet
+
+def encrypt_image(image, password):
+    # Convert the image to a NumPy array
+    image_array = np.array(image)
+
+    # Generate a key from the password
+    key = Fernet.generate_key()
+    fernet = Fernet(key)
+
+    # Encrypt the image data using the key
+    encrypted_image_data = fernet.encrypt(image_array.tobytes())
+
+    # Add Gaussian noise to the encrypted image data
+    noise = np.random.normal(0, 30, encrypted_image_data.shape).astype('uint8')
+    encrypted_image_data += noise
+
+    return encrypted_image_data, key
+
+def decrypt_image(encrypted_image_data, key, password):
+    # Check if the password is correct
+    try:
+        fernet = Fernet(key)
+    except Exception:
+        return "Incorrect password"
+
+    # Remove the Gaussian noise from the encrypted image data
+    encrypted_image_data = encrypted_image_data.astype('int16') - 30
+    encrypted_image_data = encrypted_image_data.clip(0, 255).astype('uint8')
+
+    # Decrypt the image data using the key
+    decrypted_image_data = fernet.decrypt(encrypted_image_data.tobytes())
+
+    # Convert the decrypted data back to an image
+    decrypted_image = cv2.imdecode(np.frombuffer(decrypted_image_data, np.uint8), cv2.IMREAD_COLOR)
+
+    return decrypted_image
+
+# Create the Gradio interface
+image_input = gr.inputs.Image(label="Upload Image")
+password_input = gr.inputs.Textbox(label="Enter Password")
+encrypt_button = gr.Button("Encrypt Image")
+decrypt_button = gr.Button("Decrypt Image")
+output_image = gr.outputs.Image(label="Output Image")
+
+# Define the Gradio functions
+def encrypt(image, password):
+    encrypted_data, key = encrypt_image(image, password)
+    return encrypted_data, key
+
+def decrypt(encrypted_data, key, password):
+    decrypted_image = decrypt_image(encrypted_data, key, password)
+    return decrypted_image
+
+# Create the Gradio app
+app = gr.Interface(
+    fn=encrypt,
+    inputs=[image_input, password_input, encrypt_button],
+    outputs=[output_image, gr.outputs.Textbox(label="Encryption Key")],
+    examples=None,
+)
+
+# Add the decrypt functionality to the app
+app.add(
+    fn=decrypt,
+    inputs=[gr.inputs.Image(label="Encrypted Image"), gr.inputs.Textbox(label="Encryption Key"), password_input, decrypt_button],
+    outputs=output_image,
+    examples=None,
+)
+
+# Launch the app
+app.launch()