Update app.py

app.py

@@ -5,59 +5,63 @@ import numpy as np
 import tensorflow as tf
 import gradio as gr
 
-# Load your trained model without trying to restore the old optimizer/loss,
-# then re-compile it so we can call predict()
+# Load the trained model without restoring old optimizer/loss, then re-compile
 model = tf.keras.models.load_model("best_model.h5", compile=False)
 model.compile(optimizer="adam", loss="mse")
 
 def process_and_denoise(image):
     """
     Takes any input image (color or grayscale), converts it to 64×64 grayscale,
-    adds Gaussian noise, and returns:
-      1) the resized original
-      2) the noisy version
-      3) the model’s denoised reconstruction
+    adds Gaussian noise, runs the autoencoder, and returns all three stages
+    upsampled to 128×128 for clearer display:
+      1) Original resized
+      2) Noisy version
+      3) Denoised reconstruction
     """
-    # If color, convert to gray; otherwise accept 1-channel
+    # Convert to single-channel grayscale
     if image.ndim == 3 and image.shape[2] == 3:
         gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
     else:
-        # Could be (H,W,1) or (H,W)
         gray = image[..., 0] if image.ndim == 3 else image
 
-    # Resize and normalize
-    orig = cv2.resize(gray, (64, 64))
-    orig_norm = orig.astype(np.float32) / 255.0
+    # Resize down to 64×64 for the model
+    small = cv2.resize(gray, (64, 64))
+    norm   = small.astype(np.float32) / 255.0
 
     # Add Gaussian noise
     sigma = 0.1
-    noisy = orig_norm + sigma * np.random.randn(*orig_norm.shape)
+    noisy = norm + sigma * np.random.randn(*norm.shape)
     noisy = np.clip(noisy, 0.0, 1.0)
 
-    # Denoise via the autoencoder
-    inp = noisy[np.newaxis, ..., np.newaxis]    # shape (1,64,64,1)
-    pred = model.predict(inp)[0, ..., 0]        # shape (64,64)
+    # Denoise with the autoencoder
+    inp  = noisy[np.newaxis, ..., np.newaxis]   # (1,64,64,1)
+    pred = model.predict(inp)[0, ..., 0]         # (64,64)
 
-    # Convert back to uint8 for display
-    orig_disp = (orig_norm  * 255).astype(np.uint8)
-    noisy_disp = (noisy      * 255).astype(np.uint8)
-    recon_disp = (pred       * 255).astype(np.uint8)
+    # Convert back to uint8
+    orig_disp = (norm  * 255).astype(np.uint8)
+    noisy_disp = (noisy * 255).astype(np.uint8)
+    recon_disp = (pred  * 255).astype(np.uint8)
 
-    return orig_disp, noisy_disp, recon_disp
+    # Upsample each to 128×128 for display
+    def upsample(img):
+        return cv2.resize(img, (128, 128), interpolation=cv2.INTER_LINEAR)
+
+    return upsample(orig_disp), upsample(noisy_disp), upsample(recon_disp)
 
 demo = gr.Interface(
     fn=process_and_denoise,
     inputs=gr.Image(type="numpy", label="Input Image"),
     outputs=[
-        gr.Image(type="numpy", label="Original (64×64)"),
-        gr.Image(type="numpy", label="Noisy (σ=0.1)"),
-        gr.Image(type="numpy", label="Denoised Output")
+        gr.Image(type="numpy", label="Original ▶ 128×128"),
+        gr.Image(type="numpy", label="Noisy (σ=0.1) ▶ 128×128"),
+        gr.Image(type="numpy", label="Denoised ▶ 128×128")
     ],
     title="Denoising Autoencoder Demo",
     description=(
-        "Upload any image (grayscale or color). "
-        "This will convert it to 64×64 grayscale, add Gaussian noise, "
-        "and then denoise it with the trained autoencoder."
+        "Upload any image (grayscale or color).\n"
+        "- Internally resized to 64×64 for denoising.\n"
+        "- Adds Gaussian noise (σ=0.1).\n"
+        "- Displays all stages upsampled to 128×128 for clearer viewing."
     )
 )