Update app.py

app.py

@@ -2,10 +2,15 @@ import numpy as np
 from tensorflow.keras.models import load_model
 from PIL import Image
 import gradio as gr
+from datasets import load_dataset
+import random
 
-# Load model
+# Load pre-trained model
 model = load_model("unet_model.h5", compile=False)
 
+# Load Hugging Face dataset
+dataset = load_dataset("AIOmarRehan/Cropped_Yale_Faces")
+
 # Preprocess function
 def preprocess_image(image, target_size=(192, 176)):
     image = image.resize((target_size[1], target_size[0]))  # width, height
@@ -14,34 +19,42 @@ def preprocess_image(image, target_size=(192, 176)):
         image = np.expand_dims(image, axis=-1)
     return np.expand_dims(image, axis=0)
 
-# Prediction function for Gradio
-def predict(img):
-    # Convert to grayscale
+# Prediction function
+def predict(img=None, use_dataset=False):
+    if use_dataset:
+        # Pick a random sample from the dataset
+        example = random.choice(dataset["train"])
+        img = example["image"]
+
+    if img is None:
+        return None
+
     img = img.convert("L")
 
-    # Preprocess
     input_data = preprocess_image(img)
 
-    # Model prediction
     pred = model.predict(input_data)[0]
 
-    # Remove channel
     if pred.ndim == 3 and pred.shape[-1] == 1:
         pred = np.squeeze(pred, axis=-1)
 
-    # Convert to image
     pred_img = (pred * 255).astype(np.uint8)
     pred_img = Image.fromarray(pred_img)
 
     return pred_img
 
-# Gradio UI
+# Gradio Interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Upload Image"),
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Checkbox(label="Use Random Dataset Image")
+    ],
     outputs=gr.Image(type="pil", label="Denoised Output"),
     title="U-Net Image Denoising",
-    description="Upload a grayscale image and get the denoised result using a U-Net model."
+    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."
 )
 
+# Launch the app
 interface.launch()