Update app.py

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 from datasets import load_dataset
 import random
 
-# Load pre-trained model
+# Load pre-trained U-Net model-
 model = load_model("unet_model.h5", compile=False)
 
 # Load Hugging Face dataset
@@ -19,41 +19,66 @@ def preprocess_image(image, target_size=(192, 176)):
         image = np.expand_dims(image, axis=-1)
     return np.expand_dims(image, axis=0)
 
+# Salt-and-pepper noise function
+def add_salt_and_pepper_noise(image, amount=0.05):
+    """
+    image: PIL Image in grayscale ('L') or RGB
+    amount: fraction of pixels to corrupt
+    """
+    img_array = np.array(image)
+    # Salt noise
+    num_salt = np.ceil(amount * img_array.size * 0.5)
+    coords = [np.random.randint(0, i - 1, int(num_salt)) for i in img_array.shape]
+    img_array[tuple(coords)] = 255
+    # Pepper noise
+    num_pepper = np.ceil(amount * img_array.size * 0.5)
+    coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in img_array.shape]
+    img_array[tuple(coords)] = 0
+    return Image.fromarray(img_array)
+
 # Prediction function
-def predict(img=None, use_dataset=False):
+def predict(img=None, use_dataset=False, add_noise=False):
     if use_dataset:
-        # Pick a random sample from the dataset
+        # Pick random image from dataset
         example = random.choice(dataset["train"])
         img = example["image"]
 
     if img is None:
-        return None
+        return None, None
 
     img = img.convert("L")
 
-    input_data = preprocess_image(img)
+    noisy_img = img
+    if add_noise:
+        noisy_img = add_salt_and_pepper_noise(img)
+
+    input_data = preprocess_image(noisy_img)
 
     pred = model.predict(input_data)[0]
 
     if pred.ndim == 3 and pred.shape[-1] == 1:
         pred = np.squeeze(pred, axis=-1)
 
-    pred_img = (pred * 255).astype(np.uint8)
-    pred_img = Image.fromarray(pred_img)
+    denoised_img = (pred * 255).astype(np.uint8)
+    denoised_img = Image.fromarray(denoised_img)
 
-    return pred_img
+    return noisy_img, denoised_img
 
 # Gradio Interface
 interface = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
-        gr.Checkbox(label="Use Random Dataset Image")
+        gr.Checkbox(label="Use Random Dataset Image"),
+        gr.Checkbox(label="Add Salt-and-Pepper Noise")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Noisy Input Image"),
+        gr.Image(type="pil", label="Denoised Output Image")
     ],
-    outputs=gr.Image(type="pil", label="Denoised Output"),
-    title="U-Net Image Denoising",
-    description="Upload a grayscale image or use a random image from the Cropped Yale Faces dataset. "
-                "The U-Net model will provide a denoised output."
+    title="U-Net Image Denoising with Salt-and-Pepper Noise",
+    description="Upload an image or pick a random image from the Cropped Yale Faces dataset. "
+                "Optionally add salt-and-pepper noise to the image before denoising."
 )
 
 # Launch the app