Add sample images as alternative to upload

app.py

@@ -5,6 +5,10 @@ import numpy as np
 from PIL import Image
 import math
 from einops import rearrange
+import os
+import glob
+import base64
+from io import BytesIO
 
 
 def to_2tuple(x):
@@ -722,6 +726,14 @@ def upscale_image(image):
     return Image.fromarray(output_np)
 
 
+# Get sample images
+def get_sample_images():
+    sample_dir = "sample_images"
+    if os.path.exists(sample_dir):
+        image_files = glob.glob(os.path.join(sample_dir, "*.png")) + glob.glob(os.path.join(sample_dir, "*.jpg"))
+        return sorted(image_files)
+    return []
+
 # Gradio interface using Blocks for better layout control
 def upscale_and_display(image):
     if image is None:
@@ -731,8 +743,23 @@ def upscale_and_display(image):
     upscaled = upscale_image(image)
     return upscaled
 
-# Custom CSS to make images fill containers while preserving aspect ratio
-css = """
+def select_sample_image(image_path):
+    if image_path:
+        return Image.open(image_path)
+    return None
+
+def image_to_base64(image_path):
+    """Convert image to base64 data URL for CSS background"""
+    img = Image.open(image_path)
+    img.thumbnail((120, 120), Image.Resampling.LANCZOS)
+    buffer = BytesIO()
+    img.save(buffer, format='PNG')
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/png;base64,{img_str}"
+
+# Generate CSS with base64 images
+def generate_css():
+    base_css = """
 /* Target only the image display area, not the whole component */
 .image-container [data-testid="image"] {
     height: 500px !important;
@@ -746,42 +773,81 @@ css = """
     object-fit: contain !important;
     object-position: center !important;
 }
+
+/* Sample image buttons with background images */
+.sample-image-btn {
+    height: 120px !important;
+    width: 120px !important;
+    background-size: cover !important;
+    background-position: center !important;
+    border: 2px solid #ddd !important;
+    border-radius: 8px !important;
+    cursor: pointer !important;
+    transition: border-color 0.2s !important;
+    margin: 5px !important;
+}
+
+.sample-image-btn:hover {
+    border-color: #007acc !important;
+}
 """
 
+    # Add background images for each sample
+    sample_images = get_sample_images()
+    for i, img_path in enumerate(sample_images):
+        base64_img = image_to_base64(img_path)
+        base_css += f"#sample_btn_{i} {{ background-image: url('{base64_img}'); }}\n"
+
+    return base_css
+
+css = generate_css()
+
 with gr.Blocks(css=css, title="HAT Super-Resolution for Satellite Images") as iface:
     gr.Markdown("# HAT Super-Resolution for Satellite Images")
-    gr.Markdown("Upload a satellite image to enhance its resolution by 4x using a fine-tuned HAT model. The output shows the full high-resolution enhanced image.")
+    gr.Markdown("Upload a satellite image or select a sample to enhance its resolution by 4x.")
+
+    # Sample images
+    sample_images = get_sample_images()
+    sample_buttons = []
+    if sample_images:
+        gr.Markdown("**Sample Images (click to select):**")
+        with gr.Row():
+            for i, img_path in enumerate(sample_images):
+                btn = gr.Button(
+                    "",
+                    elem_id=f"sample_btn_{i}",
+                    elem_classes="sample-image-btn"
+                )
+                sample_buttons.append((btn, img_path))
 
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(
-                type="pil",
-                label="Input Satellite Image",
-                elem_classes="image-container",
-                sources=["upload"],
-                height=500,
-                width=500,
-                show_download_button=False
-            )
-
-        with gr.Column():
-            output_image = gr.Image(
-                type="pil",
-                label="Enhanced Output (4x Super-Resolution)",
-                elem_classes="image-container",
-                interactive=False,
-                height=500,
-                width=500,
-                show_download_button=True
-            )
+        input_image = gr.Image(
+            type="pil",
+            label="Input Image",
+            elem_classes="image-container",
+            sources=["upload"],
+            height=500,
+            width=500
+        )
+
+        output_image = gr.Image(
+            type="pil",
+            label="Enhanced Output (4x)",
+            elem_classes="image-container",
+            interactive=False,
+            height=500,
+            width=500,
+            show_download_button=True
+        )
 
     submit_btn = gr.Button("Enhance Image", variant="primary")
 
-    submit_btn.click(
-        fn=upscale_and_display,
-        inputs=input_image,
-        outputs=output_image
-    )
+    # Event handlers
+    if sample_images:
+        for btn, img_path in sample_buttons:
+            btn.click(fn=lambda path=img_path: select_sample_image(path), outputs=input_image)
+
+    submit_btn.click(fn=upscale_and_display, inputs=input_image, outputs=output_image)
 
 if __name__ == "__main__":
     iface.launch()