Update app.py

app.py

@@ -1,75 +1,41 @@
-import cv2
-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.Textbox(label="Encryption Key")],
-    examples=None,
-)
-
-# Add the decrypt functionality to the app
-app.add(
-    fn=decrypt,
-    inputs=[gr.Image(label="Encrypted Image"), gr.Textbox(label="Encryption Key"), password_input, decrypt_button],
-    outputs=output_image,
-    examples=None,
+import numpy as np
+from PIL import Image
+from skimage.util import random_noise
+from skimage.restoration import denoise_nl_means, estimate_sigma
+
+def add_noise(image, mean, var):
+    # Convert variance to noise standard deviation
+    sigma = var ** 0.5
+    # Add Gaussian noise to the image
+    noisy_image = random_noise(image, mode='gaussian', mean=mean, var=var)
+    noisy_image = np.array(255 * noisy_image, dtype=np.uint8)  # Scale back to 0-255
+    return Image.fromarray(noisy_image)
+
+def remove_noise(image):
+    # Convert image to grayscale for denoising (optional depending on need)
+    image = image.convert('L')
+    image_np = np.array(image)
+    # Estimate noise sigma from image
+    sigma_est = np.mean(estimate_sigma(image_np, multichannel=False))
+    # Apply Non-Local Means Denoising
+    denoised_image = denoise_nl_means(image_np, h=1.15 * sigma_est, fast_mode=True,
+                                      patch_size=5, patch_distance=6, multichannel=False)
+    denoised_image = np.array(255 * denoised_image, dtype=np.uint8)  # Scale back to 0-255
+    return Image.fromarray(denoised_image)
+
+# Define Gradio interface
+interface = gr.Interface(
+    fn=[add_noise, remove_noise],
+    inputs=[
+        gr.Image(shape=None, label="Original Image"), 
+        gr.Slider(0, 0.05, step=0.001, default=0, label="Mean of Gaussian Noise"),
+        gr.Slider(0, 0.01, step=0.0001, default=0.001, label="Variance of Gaussian Noise")
+    ],
+    outputs=[gr.Image(label="Processed Image")],
+    title="Image Noise Addition and Removal",
+    description="Add Gaussian noise to an image and attempt to remove it. Adjust the mean and variance for different effects."
 )
 
 # Launch the app
-app.launch()
+interface.launch()