Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import os
 import torch
 import torch.nn.functional as F
@@ -203,19 +202,19 @@ def click_to_similarity_in_same_image(
     return marked_ref, heatmap_pil, overlay_pil, overlay_boxes_pil
 
 # ----------------------------
-# Gradio UI (No changes needed here)
+# Gradio UI
 # ----------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similarity") as demo:
     gr.Markdown("# 🦖 DINOv3 Single-Image Patch Similarity")
     gr.Markdown("## Running on CPU-only Space, feature extraction after uploading an image can take a moment")
-    gr.Markdown("Upload one image, then **click anywhere** to highlight the most similar regions in the *same* image.")
+    gr.Markdown("1. Upload an image. 2. Click **Process Image**. 3. **Click anywhere on the processed image** to find similar regions.")
 
     app_state = gr.State()
 
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(
-                label="Image (click anywhere)",
+                label="1. Upload Image",
                 type="pil",
                 value="https://images.squarespace-cdn.com/content/v1/607f89e638219e13eee71b1e/1684821560422-SD5V37BAG28BURTLIXUQ/michael-sum-LEpfefQf4rU-unsplash.jpg"
             )
@@ -224,6 +223,9 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similari
                 label="Processing Resolution",
                 info="Higher values = more detail but slower processing",
             )
+            # ⭐️ ADDED BUTTON
+            process_button = gr.Button("2. Process Image", variant="primary")
+            
             with gr.Row():
                 alpha = gr.Slider(0.0, 1.0, value=0.55, step=0.05, label="Overlay opacity")
                 cmap = gr.Dropdown(
@@ -236,22 +238,25 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similari
             box_radius = gr.Slider(0, 10, value=1, step=1, label="Box radius (patches)")
 
     with gr.Row():
-        marked_image = gr.Image(label="Click marker", interactive=False)
+        marked_image = gr.Image(label="3. Click on this image", interactive=True) # Changed from False to True for clarity
         heatmap_output = gr.Image(label="Similarity heatmap", interactive=False)
     with gr.Row():
         overlay_output = gr.Image(label="Overlay (image ⊕ heatmap)", interactive=False)
         overlay_boxes_output = gr.Image(label="Overlay + top-K similar patch boxes", interactive=False)
 
-    def _on_upload_or_slider_change(img: Image.Image, long_side: int, progress=gr.Progress(track_tqdm=True)):
+    # ⭐️ RENAMED FUNCTION
+    def _process_image(img: Image.Image, long_side: int, progress=gr.Progress(track_tqdm=True)):
         if img is None:
+            gr.Warning("Please upload an image first!")
             return None, None
         progress(0, desc="Extracting features...")
         st = extract_image_features(img, int(long_side))
-        progress(1, desc="Done!")
+        progress(1, desc="Done! You can now click on the image.")
         return st["img"], st
 
     def _on_click(st, a: float, m: str, excl: int, k: int, box_rad: int, evt: gr.SelectData):
         if not st or evt is None:
+            gr.Warning("Please process an image before clicking on it.")
             return None, None, None, None
         return click_to_similarity_in_same_image(
             st, click_xy=evt.index, exclude_radius_patches=int(excl),
@@ -259,14 +264,18 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Single-Image Patch Similari
             box_radius_patches=int(box_rad),
         )
 
-    # Wire events
-    inputs_for_update = [input_image, target_long_side]
-    outputs_for_update = [marked_image, app_state]
+    # ⭐️ UPDATED EVENT WIRING
+    inputs_for_processing = [input_image, target_long_side]
+    outputs_for_processing = [marked_image, app_state]
     
-    input_image.upload(_on_upload_or_slider_change, inputs=inputs_for_update, outputs=outputs_for_update)
-    target_long_side.change(_on_upload_or_slider_change, inputs=inputs_for_update, outputs=outputs_for_update)
-    demo.load(_on_upload_or_slider_change, inputs=inputs_for_update, outputs=outputs_for_update) # Process default image on load
-
+    # The button now triggers the main processing function
+    process_button.click(
+        _process_image, 
+        inputs=inputs_for_processing, 
+        outputs=outputs_for_processing
+    )
+    
+    # The click event on the image remains the same
     marked_image.select(
         _on_click,
         inputs=[app_state, alpha, cmap, exclude_r, topk, box_radius],