Update app.py

app.py

@@ -22,9 +22,10 @@ def add_noise(image, password):
     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."""
-    # Check if the noisy_image is a numpy array and convert it to PIL Image if so
     if isinstance(noisy_image, np.ndarray):
         noisy_image = Image.fromarray(noisy_image.astype('uint8'))
 
@@ -32,15 +33,26 @@ def remove_noise(noisy_image, password):
     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
+    image = noisy_image.convert('L')  # Ensure image is in grayscale for simplicity
     image_np = np.array(image)
-    sigma_est = np.mean(estimate_sigma(image_np, multichannel=False))
-    denoised_image = denoise_nl_means(image_np, h=1.15 * sigma_est, fast_mode=True,
+
+    # Convert image to floating point type needed for denoise_nl_means
+    image_float = img_as_float(image_np)
+
+    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))
+
+    denoised_image = denoise_nl_means(image_float, h=1.15 * sigma_est, fast_mode=True,
                                       patch_size=5, patch_distance=6, multichannel=False)
     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")