Update app.py

app.py

@@ -1,4 +1,5 @@
 # app.py
+
 import os
 import torch
 import torch.nn.functional as F
@@ -6,47 +7,79 @@ import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
 import torchvision.transforms.functional as TF
-from matplotlib import colormaps
-from transformers import AutoModel
+
+# --- Robust colormap import (Matplotlib ≥3.5 and older versions) ---
+try:
+    from matplotlib import colormaps as _mpl_colormaps
+    def _get_cmap(name: str):
+        return _mpl_colormaps[name]
+except Exception:
+    import matplotlib.cm as _cm
+    def _get_cmap(name: str):
+        return _cm.get_cmap(name)
+
+from transformers import AutoModel  # uses trust_remote_code for DINOv3
 
 # ----------------------------
 # Configuration
 # ----------------------------
-MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
+# Default to smaller/faster ViT-S/16+; offer ViT-H/16+ as alternative.
+DEFAULT_MODEL_ID = "facebook/dinov3-vits16plus-pretrain-lvd1689m"
+ALT_MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
+AVAILABLE_MODELS = [DEFAULT_MODEL_ID, ALT_MODEL_ID]
+
 PATCH_SIZE = 16
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Normalization constants (standard for ImageNet)
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
 
 # ----------------------------
-# Model Loading (Hugging Face Hub)
+# Model Loading (Hugging Face Hub) with caching
 # ----------------------------
-def load_model_from_hub():
-    """Loads the DINOv3 model from the Hugging Face Hub."""
-    print(f"Loading model '{MODEL_ID}' from Hugging Face Hub...")
+_model_cache = {}
+_current_model_id = None
+model = None  # global reference used by extract_image_features()
+
+def load_model_from_hub(model_id: str):
+    """Loads a DINOv3 model from the Hugging Face Hub."""
+    print(f"Loading model '{model_id}' from Hugging Face Hub...")
     try:
-        token = os.environ.get("HF_TOKEN")
-        model = AutoModel.from_pretrained(MODEL_ID, token=token, trust_remote_code=True)
-        model.to(DEVICE).eval()
-        print(f"✅ Model loaded successfully on device: {DEVICE}")
-        return model
+        token = os.environ.get("HF_TOKEN")  # optional, for gated models
+        mdl = AutoModel.from_pretrained(model_id, token=token, trust_remote_code=True)
+        mdl.to(DEVICE).eval()
+        print(f"✅ Model '{model_id}' loaded successfully on device: {DEVICE}")
+        return mdl
     except Exception as e:
-        print(f"❌ Failed to load model: {e}")
+        print(f"❌ Failed to load model '{model_id}': {e}")
         raise gr.Error(
-            f"Could not load model '{MODEL_ID}'. "
-            "This is a gated model. Please ensure you have accepted the terms on its Hugging Face page "
-            "and set your HF_TOKEN as a secret in your Space settings. "
+            f"Could not load model '{model_id}'. "
+            "If the model is gated, please accept the terms on its Hugging Face page "
+            "and set HF_TOKEN in your environment. "
             f"Original error: {e}"
         )
 
-# Load the model globally when the app starts
-model = load_model_from_hub()
+def get_model(model_id: str):
+    """Return a cached model if available, otherwise load and cache it."""
+    if model_id in _model_cache:
+        return _model_cache[model_id]
+    mdl = load_model_from_hub(model_id)
+    _model_cache[model_id] = mdl
+    return mdl
+
+# Load default model at startup
+model = get_model(DEFAULT_MODEL_ID)
+_current_model_id = DEFAULT_MODEL_ID
 
 # ----------------------------
 # Helper Functions (resize, viz)
 # ----------------------------
 def resize_to_grid(img: Image.Image, long_side: int, patch: int) -> torch.Tensor:
+    """
+    Resizes so max(h,w)=long_side (keeping aspect), then rounds each side UP to a multiple of 'patch'.
+    Returns CHW float tensor in [0,1].
+    """
     w, h = img.size
     scale = long_side / max(h, w)
     new_h = max(patch, int(round(h * scale)))
@@ -58,13 +91,17 @@ def resize_to_grid(img: Image.Image, long_side: int, patch: int) -> torch.Tensor
 def colorize(sim_map_up: np.ndarray, cmap_name: str = "viridis") -> Image.Image:
     x = sim_map_up.astype(np.float32)
     x = (x - x.min()) / (x.max() - x.min() + 1e-6)
-    rgb = (colormaps[cmap_name](x)[..., :3] * 255).astype(np.uint8)
+    rgb = (_get_cmap(cmap_name)(x)[..., :3] * 255).astype(np.uint8)
     return Image.fromarray(rgb)
 
 def blend(base: Image.Image, heat: Image.Image, alpha: float = 0.55) -> Image.Image:
+    # Put alpha on heatmap and composite for a crisp overlay
     base = base.convert("RGBA")
     heat = heat.convert("RGBA")
-    return Image.blend(base, heat, alpha=alpha)
+    a = Image.new("L", heat.size, int(255 * alpha))
+    heat.putalpha(a)
+    out = Image.alpha_composite(base, heat)
+    return out.convert("RGB")
 
 def draw_crosshair(img: Image.Image, x: int, y: int, radius: int = None) -> Image.Image:
     r = radius if radius is not None else max(2, PATCH_SIZE // 2)
@@ -96,26 +133,33 @@ def patch_neighborhood_box(r: int, c: int, Hp: int, Wp: int, rad: int, patch: in
     return (x0, y0, x1, y1)
 
 # ----------------------------
-# Feature Extraction
+# Feature Extraction (using transformers)
 # ----------------------------
 @torch.inference_mode()
 def extract_image_features(image_pil: Image.Image, target_long_side: int):
+    """
+    Extracts patch features from an image using the loaded Hugging Face model.
+    """
     t = resize_to_grid(image_pil, target_long_side, PATCH_SIZE)
     t_norm = TF.normalize(t, IMAGENET_MEAN, IMAGENET_STD).unsqueeze(0).to(DEVICE)
     _, _, H, W = t_norm.shape
     Hp, Wp = H // PATCH_SIZE, W // PATCH_SIZE
-    
+
+    # Models output: [CLS] + 4 register tokens + patches
     outputs = model(t_norm)
-    
+
+    # Skip the 5 special tokens to get only patch embeddings
     n_special_tokens = 5
     patch_embeddings = outputs.last_hidden_state.squeeze(0)[n_special_tokens:, :]
-    
+
+    # L2-normalize features for cosine similarity
     X = F.normalize(patch_embeddings, p=2, dim=-1)
+
     img_resized = TF.to_pil_image(t)
     return {"X": X, "Hp": Hp, "Wp": Wp, "img": img_resized}
 
 # ----------------------------
-# Similarity Logic
+# Similarity inside the same image
 # ----------------------------
 def click_to_similarity_in_same_image(
     state: dict,
@@ -128,17 +172,21 @@ def click_to_similarity_in_same_image(
 ):
     if not state:
         return None, None, None, None
+
     X = state["X"]
     Hp, Wp = state["Hp"], state["Wp"]
     base_img = state["img"]
     img_w, img_h = base_img.size
+
     x_pix, y_pix = click_xy
     col = int(np.clip(x_pix // PATCH_SIZE, 0, Wp - 1))
     row = int(np.clip(y_pix // PATCH_SIZE, 0, Hp - 1))
     idx = row * Wp + col
+
     q = X[idx]
     sims = torch.matmul(X, q)
     sim_map = sims.view(Hp, Wp)
+
     if exclude_radius_patches > 0:
         rr, cc = torch.meshgrid(
             torch.arange(Hp, device=sims.device),
@@ -147,14 +195,17 @@ def click_to_similarity_in_same_image(
         )
         mask = (torch.abs(rr - row) <= exclude_radius_patches) & (torch.abs(cc - col) <= exclude_radius_patches)
         sim_map = sim_map.masked_fill(mask, float("-inf"))
+
     sim_up = F.interpolate(
         sim_map.unsqueeze(0).unsqueeze(0),
         size=(img_h, img_w),
         mode="bicubic",
         align_corners=False,
     ).squeeze().detach().cpu().numpy()
+
     heatmap_pil = colorize(sim_up, cmap_name)
     overlay_pil = blend(base_img, heatmap_pil, alpha=alpha)
+
     overlay_boxes_pil = overlay_pil
     if topk and topk > 0:
         flat = sim_map.view(-1)
@@ -171,83 +222,100 @@ def click_to_similarity_in_same_image(
                 for r, c in [divmod(j.item(), Wp) for j in top_idx]
             ]
             overlay_boxes_pil = draw_boxes(overlay_pil, boxes, outline="yellow", width=3, labels=True)
+
     marked_ref = draw_crosshair(base_img, x_pix, y_pix, radius=PATCH_SIZE // 2)
     return marked_ref, heatmap_pil, overlay_pil, overlay_boxes_pil
 
 # ----------------------------
-# Gradio UI
+# Gradio UI (+ Start button, + Model dropdown)
 # ----------------------------
-with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Patch Similarity") as demo:
-    gr.Markdown("# 🦖 DINOv3: Visualizing Patch Similarity")
-    gr.Markdown(
-        "Upload an image, then **click anywhere** on it to find the most visually similar regions. "
-        "**Note:** If running on a CPU-only Space, feature extraction after uploading an image can take a moment."
-    )
-    
+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.")
+
     app_state = gr.State()
-    
+
     with gr.Row():
-        with gr.Column(scale=2):
+        with gr.Column(scale=1):
             input_image = gr.Image(
                 label="Image (click anywhere)",
                 type="pil",
                 value="https://images.squarespace-cdn.com/content/v1/607f89e638219e13eee71b1e/1684821560422-SD5V37BAG28BURTLIXUQ/michael-sum-LEpfefQf4rU-unsplash.jpg"
             )
-            with gr.Accordion("⚙️ Visualization Controls", open=True):
-                target_long_side = gr.Slider(
-                    minimum=224, maximum=1024, value=768, step=16,
-                    label="Processing Resolution",
-                    info="Higher values = more detail but slower processing",
-                )
-                alpha = gr.Slider(0.0, 1.0, value=0.55, step=0.05, label="Overlay Opacity")
+            target_long_side = gr.Slider(
+                minimum=224, maximum=1024, value=768, step=16,
+                label="Processing Resolution",
+                info="Higher values = more detail but slower processing",
+            )
+            with gr.Row():
+                alpha = gr.Slider(0.0, 1.0, value=0.55, step=0.05, label="Overlay opacity")
                 cmap = gr.Dropdown(
                     ["viridis", "magma", "plasma", "inferno", "turbo", "cividis"],
-                    value="viridis", label="Heatmap Colormap",
+                    value="viridis", label="Colormap",
                 )
-            with gr.Accordion("⚙️ Similarity Controls", open=True):
-                exclude_r = gr.Slider(0, 10, value=0, step=1, label="Exclude Radius (patches)", info="Ignore patches around the click point.")
-                topk = gr.Slider(0, 50, value=10, step=1, label="Top-K Boxes", info="Number of similar regions to highlight.")
-                box_radius = gr.Slider(0, 10, value=1, step=1, label="Box Radius (patches)", info="Size of the highlight box.")
-
-        with gr.Column(scale=3):
-            marked_image = gr.Image(label="Your Click (on processed image)", interactive=False)
-            with gr.Tabs():
-                with gr.TabItem("📦 Bounding Boxes"):
-                    overlay_boxes_output = gr.Image(label="Overlay + Top-K Similar Patches", interactive=False)
-                with gr.TabItem("🔥 Heatmap"):
-                    heatmap_output = gr.Image(label="Similarity Heatmap", interactive=False)
-                with gr.TabItem(" blended"):
-                    overlay_output = gr.Image(label="Blended Overlay (Image + Heatmap)", interactive=False)
-
-    def _on_upload_or_slider_change(img: Image.Image, long_side: int, progress=gr.Progress(track_tqdm=True)):
+            # NEW: Backbone selector (default = smaller/faster ViT-S/16+)
+            model_choice = gr.Dropdown(
+                choices=AVAILABLE_MODELS,
+                value=DEFAULT_MODEL_ID,
+                label="Backbone (DINOv3)",
+                info="ViT-S/16+ is smaller & faster; ViT-H/16+ is larger.",
+            )
+            # Start processing button
+            with gr.Row():
+                start_btn = gr.Button("▶️ Start processing", variant="primary")
+
+        with gr.Column(scale=1):
+            exclude_r = gr.Slider(0, 10, value=0, step=1, label="Exclude radius (patches)")
+            topk = gr.Slider(0, 200, value=20, step=1, label="Top-K boxes")
+            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)
+        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 _ensure_model(model_id: str):
+        """Ensure the global 'model' matches the dropdown selection."""
+        global model, _current_model_id
+        if model_id != _current_model_id:
+            model = get_model(model_id)
+            _current_model_id = model_id
+
+    def _on_upload_or_slider_change(img: Image.Image, long_side: int, model_id: str, progress=gr.Progress(track_tqdm=True)):
         if img is None:
             return None, None
-        progress(0, desc="🦖 Extracting DINOv3 features...")
+        _ensure_model(model_id)
+        progress(0, desc="Extracting features...")
         st = extract_image_features(img, int(long_side))
-        progress(1, desc="✅ Done!")
-        # Clear old results when a new image is uploaded
-        return st["img"], st, None, None, None, None
+        progress(1, desc="Done!")
+        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:
-            # Return current state if no click data
-            return st.get("img"), None, None, None
-        
-        marked, heat, overlay, boxes = click_to_similarity_in_same_image(
+            return None, None, None, None
+        return click_to_similarity_in_same_image(
             st, click_xy=evt.index, exclude_radius_patches=int(excl),
             topk=int(k), alpha=float(a), cmap_name=m,
             box_radius_patches=int(box_rad),
         )
-        return marked, heat, overlay, boxes
 
     # Wire events
-    inputs_for_update = [input_image, target_long_side]
-    outputs_for_upload = [marked_image, app_state, heatmap_output, overlay_output, overlay_boxes_output, marked_image]
+    inputs_for_update = [input_image, target_long_side, model_choice]
+    outputs_for_update = [marked_image, app_state]
+
+    # Auto triggers (kept)
+    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)
+    model_choice.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
 
-    input_image.upload(_on_upload_or_slider_change, inputs=inputs_for_update, outputs=outputs_for_upload)
-    target_long_side.change(_on_upload_or_slider_change, inputs=inputs_for_update, outputs=outputs_for_upload)
-    demo.load(_on_upload_or_slider_change, inputs=inputs_for_update, outputs=outputs_for_upload)
+    # Manual trigger via button (kept)
+    start_btn.click(_on_upload_or_slider_change, inputs=inputs_for_update, outputs=outputs_for_update)
 
+    # Click to compute similarities
     marked_image.select(
         _on_click,
         inputs=[app_state, alpha, cmap, exclude_r, topk, box_radius],
@@ -255,4 +323,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Patch Similarity") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()