Upload app.py

app.py

@@ -1,3 +1,4 @@
+import glob
 import gradio as gr
 import torch
 from transformers import ViTForImageClassification, AutoImageProcessor
@@ -5,36 +6,36 @@ from PIL import Image, ImageDraw, ImageFont
 import random
 import os
 
-# Load model and processor from Hugging Face
+# ✅ Load model and processor from Hugging Face
 model_name = "akridge/noaa-esd-coral-bleaching-vit-classifier-v1"
 model = ViTForImageClassification.from_pretrained(model_name)
 processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
 
-# Ensure id2label keys are integers and labels are uppercase
+# 🗂️ Ensure id2label keys are integers and labels are uppercase
 id2label = {int(k): v.upper() for k, v in model.config.id2label.items()}
 
-# Label colors (RGBA)
+# 🎨 Label colors (RGBA)
 LABEL_COLORS = {
     "CORAL": ((0, 0, 255, 80), (0, 0, 200)),        # Fill: Blue transparent, Border: Dark blue
-    "CORAL_BL": ((255, 255, 255, 150), (150, 150, 150))  # Fill: White transparent, Border: Gray
+    "CORAL_BL": ((255, 255, 255, 150), (150, 150, 150)),  # Fill: White transparent, Border: Gray
 }
 
 def predict_and_overlay(image, rows=2, cols=5, patch_size=224):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device).eval()
 
-    # Load image
+    # ✅ Load image
     image = image.convert('RGB')
     width, height = image.size
     scale_factor = max(width, height) / 800
     font_size = max(12, int(scale_factor * 8))
     border_width = max(3, int(scale_factor * 3))
 
-    # Create overlay
+    # 🖌️ Create overlay
     overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
     overlay_draw = ImageDraw.Draw(overlay)
 
-    # Generate sampled points
+    # ✅ Generate sampled points
     cell_width, cell_height = width / cols, height / rows
     sampled_points = [
         (random.randint(int(col * cell_width), int((col + 1) * cell_width - 1)),
@@ -44,18 +45,18 @@ def predict_and_overlay(image, rows=2, cols=5, patch_size=224):
 
     predictions = []
 
-    # Load font
+    # ✅ Load font
     try:
         font = ImageFont.truetype("arial.ttf", size=font_size)
     except IOError:
         font = ImageFont.load_default()
 
-    # Predict and draw patches
+    # 🔍 Predict and draw patches
     for x, y in sampled_points:
         left, upper = max(0, x - patch_size // 2), max(0, y - patch_size // 2)
         right, lower = min(width, left + patch_size), min(height, upper + patch_size)
 
-        # Predict label
+        # 🧠 Predict label
         patch = image.crop((left, upper, right, lower))
         inputs = processor(images=patch, return_tensors="pt").to(device)
         with torch.no_grad():
@@ -63,10 +64,10 @@ def predict_and_overlay(image, rows=2, cols=5, patch_size=224):
             pred_label = id2label.get(pred_id, "UNKNOWN")
             predictions.append(pred_label)
 
-        # Fill and border colors
+        # 🎨 Fill and border colors
         fill_color, border_color = LABEL_COLORS.get(pred_label, ((200, 200, 200, 100), (100, 100, 100)))
 
-        # Draw filled rectangle and label
+        # 🟦 Draw filled rectangle and label
         overlay_draw.rectangle([(left, upper), (right, lower)], fill=fill_color, outline=border_color, width=border_width)
 
         label_text = pred_label
@@ -77,37 +78,74 @@ def predict_and_overlay(image, rows=2, cols=5, patch_size=224):
         overlay_draw.rectangle(text_bg_coords, fill=(0, 0, 0, 200))
         overlay_draw.text((left + 3, upper - text_height - 4), label_text, fill="white", font=font)
 
-    # Merge overlay with original
+    # 🖼️ Merge overlay with original
     final_image = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
 
-    return final_image, predictions
-
-# Function to load example images
-def load_example_image(example_image):
-    return Image.open(example_image)
+    return final_image, ", ".join(predictions)
 
-# List example images
-example_images = [os.path.join("example_images", img) for img in os.listdir("coral_images")]
+# Load example images from coral_images folder
+example_images = glob.glob("coral_images/*.[jp][pn]g")
 
-# Gradio interface
+# 🚀 Gradio interface
 def gradio_interface(image, rows=2, cols=5):
+    if image is None:
+        return None, "No image uploaded. Please upload an image or select from examples."
     final_image, predictions = predict_and_overlay(image, rows, cols)
-    return final_image, ", ".join(predictions)
+    return final_image, predictions
 
-iface = gr.Interface(
-    fn=gradio_interface,
-    inputs=[
-        gr.inputs.Image(type="pil", label="Upload Coral Image", optional=True),
-        gr.inputs.Dropdown(choices=example_images, label="Or Select an Example Image"),
-        gr.inputs.Slider(1, 10, value=2, step=1, label="Rows of Sample Points"),
-        gr.inputs.Slider(1, 10, value=5, step=1, label="Columns of Sample Points"),
-    ],
-    outputs=[
-        gr.outputs.Image(type="pil", label="Image with Predictions"),
-        gr.outputs.Textbox(label="Predictions")
-    ],
-    title="NOAA ESD Coral Bleaching ViT Classifier",
-    description="Upload an image or select an example to sample points/patches and predict coral bleaching using the ViT classifier model hosted on Hugging Face."
-)
-
-iface.launch()
+app_title = "🌊 NOAA ESD Coral Bleaching ViT Classifier"
+app_description = """
+Upload a coral image or select from example images to sample points and predict coral bleaching using the ViT classifier model hosted on Hugging Face.
+
+**Model:** [akridge/noaa-esd-coral-bleaching-vit-classifier-v1](https://huggingface.co/akridge/noaa-esd-coral-bleaching-vit-classifier-v1)
+"""
+
+# Custom CSS for improved styling
+custom_css = """
+    .gradio-container h1 {
+        font-size: 2.2em;
+        text-align: center;
+    }
+    .gradio-container p {
+        font-size: 1.2em;
+    }
+    .gradio-container .gr-button {
+        font-size: 1.2em;
+    }
+"""
+
+with gr.Blocks(theme=gr.themes.Ocean(), css=custom_css, title=app_title) as interface:
+    gr.Markdown(f"<h1>{app_title}</h1>")
+    gr.Markdown(app_description)
+
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Coral Image")
+        result_output = gr.Image(type="pil", label="Predicted Results")
+
+    with gr.Row():
+        rows_slider = gr.Slider(1, 10, value=2, label="Rows of Sample Points")
+        cols_slider = gr.Slider(1, 10, value=5, label="Columns of Sample Points")
+
+    with gr.Row():
+        run_button = gr.Button("Run Prediction", variant="primary")
+        clear_button = gr.Button("Clear")
+
+    # Add example images section
+    gr.Examples(
+        examples=[[img] for img in example_images],
+        inputs=[image_input],
+        outputs=[result_output],
+        examples_per_page=6,
+        label="Example Coral Images"
+    )
+
+    # Define button actions
+    run_button.click(
+        fn=gradio_interface,
+        inputs=[image_input, rows_slider, cols_slider],
+        outputs=[result_output, gr.Textbox(label="Predictions")]
+    )
+
+    clear_button.click(lambda: (None, ""), outputs=[image_input, result_output])
+
+interface.launch()