Update app.py

app.py

@@ -17,14 +17,10 @@ def add_noise(image, password):
     np.random.seed(seed)
     mean = np.random.uniform(-0.1, 0.1)  # Wider range for mean
     var = np.random.uniform(0.05, 0.1)   # Higher variance for stronger noise
-    sigma = var ** 0.5
-    noisy_image = random_noise(np.array(image), mode='gaussian', mean=mean, var=var)
+    noisy_image = random_noise(np.array(image), mode='gaussian', mean=mean, var=var, clip=True)
     noisy_image = (255 * noisy_image).astype(np.uint8)  # Scale back to 0-255
     return Image.fromarray(noisy_image)
 
-
-from skimage import img_as_float
-
 def remove_noise(noisy_image, password):
     """Attempt to remove noise from the image using the same password."""
     if isinstance(noisy_image, np.ndarray):
@@ -32,36 +28,26 @@ def remove_noise(noisy_image, password):
 
     seed = hash_password(password) % (2**32)
     np.random.seed(seed)  # Reset the seed to generate the same noise parameters
-    mean = np.random.uniform(-0.05, 0.05)
-    var = np.random.uniform(0.001, 0.01)
-    image = noisy_image.convert('L')  # Ensure image is in grayscale for simplicity
-    image_np = np.array(image)
+    image = np.array(noisy_image)
 
     # Convert image to floating point type needed for denoise_nl_means
-    image_float = img_as_float(image_np)
+    image_float = img_as_float(image)
 
-    try:
-        # For newer skimage versions or if the image is colored
-        sigma_est = np.mean(estimate_sigma(image_float, channel_axis=None))
-    except TypeError:
-        # For older skimage versions
-        sigma_est = np.mean(estimate_sigma(image_float))
+    sigma_est = np.mean(estimate_sigma(image_float, channel_axis=-1))  # Use channel_axis=-1 for color images
 
     denoised_image = denoise_nl_means(image_float, h=1.15 * sigma_est, fast_mode=True,
-                                      patch_size=5, patch_distance=6, channel_axis=None)
+                                      patch_size=5, patch_distance=6, channel_axis=-1)
 
     denoised_image = (255 * denoised_image).astype(np.uint8)  # Scale back to 0-255
     return Image.fromarray(denoised_image)
 
-
-
 # Define Gradio interface with tabs for adding and removing noise
 with gr.Blocks() as interface:
     gr.Markdown("### Image Noise Encryption and Decryption App")
     with gr.Tab("Encrypt"):
         with gr.Row():
             image_input = gr.Image(label="Original Image")
-            password_input = gr.Textbox(label="Password for Encryption", placeholder="Enter a strong password", type="password")
+            password_input = gr.Textbox(label="Password for Encryption", type="password")
             encrypt_button = gr.Button("Encrypt Image")
         image_output = gr.Image(label="Encrypted Image")
         encrypt_button.click(add_noise, inputs=[image_input, password_input], outputs=image_output)
@@ -69,7 +55,7 @@ with gr.Blocks() as interface:
     with gr.Tab("Decrypt"):
         with gr.Row():
             image_input_decrypt = gr.Image(label="Encrypted Image")
-            password_input_decrypt = gr.Textbox(label="Password for Decryption", placeholder="Enter the same password used for encryption", type="password")
+            password_input_decrypt = gr.Textbox(label="Password for Decryption", type="password")
             decrypt_button = gr.Button("Decrypt Image")
         image_output_decrypt = gr.Image(label="Decrypted Image")
         decrypt_button.click(remove_noise, inputs=[image_input_decrypt, password_input_decrypt], outputs=image_output_decrypt)