Update app.py

app.py

@@ -1,85 +1,74 @@
-# app.py
+# app.py — DINOv3 two‑image patch‑similarity (click on image 1 → show similarities on both images)
+# Derived from: https://huggingface.co/spaces/sayedM/DINOv3-features (single‑image version)
+# Adds: second image input, dual feature extraction, and cross‑image similarity heatmaps/overlays.
 
 import os
-import torch
-import torch.nn.functional as F
-import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
-import torchvision.transforms.functional as TF
 
-# --- 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)
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from transformers import AutoModel  # trust_remote_code=True (DINOv3 on HF)
 
-from transformers import AutoModel  # uses trust_remote_code for DINOv3
+import gradio as gr
 
 # ----------------------------
-# Configuration
+# Config
 # ----------------------------
-# 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)
+IMAGENET_STD  = (0.229, 0.224, 0.225)
+N_SPECIAL_TOKENS = 5  # [CLS] + 4 registers in dinov3 HF ports
+
+# --- robust colormap import (matplotlib new/old)
+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)
 
 # ----------------------------
-# Model Loading (Hugging Face Hub) with caching
+# Model loading / cache
 # ----------------------------
 _model_cache = {}
 _current_model_id = None
-model = None  # global reference used by extract_image_features()
+model = None
 
 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")  # 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 '{model_id}': {e}")
-        raise gr.Error(
-            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}"
-        )
+    print(f"Loading model '{model_id}' from HF Hub…")
+    token = os.environ.get("HF_TOKEN")
+    mdl = AutoModel.from_pretrained(model_id, token=token, trust_remote_code=True)
+    mdl.to(DEVICE).eval()
+    print(f"✅ Loaded '{model_id}' on {DEVICE}")
+    return mdl
 
 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
+# Load default at startup
 model = get_model(DEFAULT_MODEL_ID)
 _current_model_id = DEFAULT_MODEL_ID
 
 # ----------------------------
-# Helper Functions (resize, viz)
+# Helpers
 # ----------------------------
-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].
-    """
+
+def resize_to_grid(img: Image.Image, long_side: int, patch: int = PATCH_SIZE) -> torch.Tensor:
+    """Resize so max(h,w)=long_side with aspect kept; then pad up to multiples of patch.
+    Return CHW float tensor in [0,1]."""
     w, h = img.size
     scale = long_side / max(h, w)
     new_h = max(patch, int(round(h * scale)))
@@ -95,7 +84,6 @@ def colorize(sim_map_up: np.ndarray, cmap_name: str = "viridis") -> Image.Image:
     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")
     a = Image.new("L", heat.size, int(255 * alpha))
@@ -111,80 +99,45 @@ def draw_crosshair(img: Image.Image, x: int, y: int, radius: int = None) -> Imag
     draw.line([(x, y - r), (x, y + r)], fill="red", width=3)
     return out
 
-def draw_boxes(img: Image.Image, boxes, outline="yellow", width=3, labels=True):
-    out = img.copy()
-    draw = ImageDraw.Draw(out)
-    for i, (x0, y0, x1, y1) in enumerate(boxes, start=1):
-        draw.rectangle([x0, y0, x1, y1], outline=outline, width=width)
-        if labels:
-            tx, ty = x0 + 2, y0 + 2
-            draw.text((tx, ty), str(i), fill=outline)
-    return out
-
-def patch_neighborhood_box(r: int, c: int, Hp: int, Wp: int, rad: int, patch: int = PATCH_SIZE):
-    r0 = max(0, r - rad)
-    r1 = min(Hp - 1, r + rad)
-    c0 = max(0, c - rad)
-    c1 = min(Wp - 1, c + rad)
-    x0 = int(c0 * patch)
-    y0 = int(r0 * patch)
-    x1 = int((c1 + 1) * patch) - 1
-    y1 = int((r1 + 1) * patch) - 1
-    return (x0, y0, x1, y1)
-
 # ----------------------------
-# Feature Extraction (using transformers)
+# Feature extraction
 # ----------------------------
 @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.
-    """
+def extract_image_features(image_pil: Image.Image, target_long_side: int, mdl=None):
+    global model
+    mdl = mdl or 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)
-
+    outputs = mdl(t_norm)
+    patch_emb = outputs.last_hidden_state.squeeze(0)[N_SPECIAL_TOKENS:, :]  # skip special tokens
+    X = F.normalize(patch_emb, p=2, dim=-1)  # (Hp*Wp, D), L2 norm for cosine
     img_resized = TF.to_pil_image(t)
+
     return {"X": X, "Hp": Hp, "Wp": Wp, "img": img_resized}
 
 # ----------------------------
-# Similarity inside the same image
+# Similarity utilities
 # ----------------------------
-def click_to_similarity_in_same_image(
-    state: dict,
-    click_xy: tuple[int, int],
-    exclude_radius_patches: int = 1,
-    topk: int = 10,
-    alpha: float = 0.55,
-    cmap_name: str = "viridis",
-    box_radius_patches: int = 4,
-):
-    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
+
+def index_from_xy(x_pix: int, y_pix: int, Wp: int) -> int:
+    col = int(np.clip(x_pix // PATCH_SIZE, 0, Wp - 1))
+    row = int(np.clip(y_pix // PATCH_SIZE, 0, (x_pix*0 + y_pix) // PATCH_SIZE))  # placeholder row calc replaced below
+    return row * Wp + col
+
+# Corrected row/col computation helper
+
+def row_col_from_xy(x_pix: int, y_pix: int, Hp: int, Wp: int):
     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
+    return row, col
 
-    q = X[idx]
-    sims = torch.matmul(X, q)
+@torch.inference_mode()
+def similarity_map(X: torch.Tensor, Hp: int, Wp: int, q_vec: torch.Tensor,
+                   img_h: int, img_w: int, exclude_radius_patches: int = 1):
+    sims = torch.matmul(X, q_vec)  # (Hp*Wp)
     sim_map = sims.view(Hp, Wp)
 
     if exclude_radius_patches > 0:
@@ -193,8 +146,8 @@ def click_to_similarity_in_same_image(
             torch.arange(Wp, device=sims.device),
             indexing="ij",
         )
-        mask = (torch.abs(rr - row) <= exclude_radius_patches) & (torch.abs(cc - col) <= exclude_radius_patches)
-        sim_map = sim_map.masked_fill(mask, float("-inf"))
+        # We'll mask later at the click location per-image if needed
+        mask_template = (rr * 0)  # kept for API parity
 
     sim_up = F.interpolate(
         sim_map.unsqueeze(0).unsqueeze(0),
@@ -202,134 +155,147 @@ def click_to_similarity_in_same_image(
         mode="bicubic",
         align_corners=False,
     ).squeeze().detach().cpu().numpy()
+    return sim_map, sim_up
+
+# ----------------------------
+# Core: click on image 1 → heatmaps on image 1 and image 2
+# ----------------------------
+
+def click_two_image_similarity(state1: dict, state2: dict, click_xy: tuple[int, int],
+                               exclude_radius_patches: int, alpha: float, cmap_name: str):
+    if not state1 or not state2:
+        return (None,)*6
+
+    X1, Hp1, Wp1, img1 = state1["X"], state1["Hp"], state1["Wp"], state1["img"]
+    X2, Hp2, Wp2, img2 = state2["X"], state2["Hp"], state2["Wp"], state2["img"]
 
-    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)
-        valid = torch.isfinite(flat)
-        if valid.any():
-            vals = flat.clone()
-            vals[~valid] = -1e9
-            k = min(topk, int(valid.sum().item()))
-            _, top_idx = torch.topk(vals, k=k, largest=True, sorted=True)
-            boxes = [
-                patch_neighborhood_box(
-                    r, c, Hp, Wp, rad=int(box_radius_patches), patch=PATCH_SIZE
-                )
-                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
+    img1_w, img1_h = img1.size
+    img2_w, img2_h = img2.size
+
+    # Build query vector from clicked patch on image 1
+    col = int(np.clip(click_xy[0] // PATCH_SIZE, 0, Wp1 - 1))
+    row = int(np.clip(click_xy[1] // PATCH_SIZE, 0, Hp1 - 1))
+    idx = row * Wp1 + col
+    q = X1[idx]  # (D,)
+
+    # Similarity on image 1
+    sims1 = torch.matmul(X1, q)
+    sim_map1 = sims1.view(Hp1, Wp1)
+    if exclude_radius_patches > 0:
+        rr, cc = torch.meshgrid(
+            torch.arange(Hp1, device=sims1.device),
+            torch.arange(Wp1, device=sims1.device),
+            indexing="ij",
+        )
+        mask1 = (torch.abs(rr - row) <= exclude_radius_patches) & (torch.abs(cc - col) <= exclude_radius_patches)
+        sim_map1 = sim_map1.masked_fill(mask1, float("-inf"))
+
+    sim1_up = F.interpolate(
+        sim_map1.unsqueeze(0).unsqueeze(0),
+        size=(img1_h, img1_w),
+        mode="bicubic",
+        align_corners=False,
+    ).squeeze().detach().cpu().numpy()
+
+    heat1 = colorize(sim1_up, cmap_name)
+    overlay1 = blend(img1, heat1, alpha)
+    marked1 = draw_crosshair(img1, int(click_xy[0]), int(click_xy[1]), radius=PATCH_SIZE // 2)
+
+    # Similarity on image 2 (no exclusion mask, since click is on image 1)
+    sims2 = torch.matmul(X2, q)
+    sim_map2 = sims2.view(Hp2, Wp2)
+    sim2_up = F.interpolate(
+        sim_map2.unsqueeze(0).unsqueeze(0),
+        size=(img2_h, img2_w),
+        mode="bicubic",
+        align_corners=False,
+    ).squeeze().detach().cpu().numpy()
+
+    heat2 = colorize(sim2_up, cmap_name)
+    overlay2 = blend(img2, heat2, alpha)
+
+    return marked1, heat1, overlay1, heat2, overlay2, sim2_up.max().item()
 
 # ----------------------------
-# Gradio UI (Manual-only processing)
+# 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("Upload one image, adjust settings, then press **▶️ Start processing**. Click on the processed image to find similar regions.")
+with gr.Blocks(theme=gr.themes.Soft(), title="DINOv3 Two‑Image Patch Similarity") as demo:
+    gr.Markdown("# DINOv3 Two‑Image Patch Similarity")
+    gr.Markdown(
+        "Upload two images and press **▶️ Process**. Click a location on **Image 1** to see similar regions on **both** images."
+    )
 
-    app_state = gr.State()
+    state1 = gr.State()
+    state2 = gr.State()
 
     with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(
-                label="Image (click anywhere after processing)",
-                type="pil",
-                value="https://images.squarespace-cdn.com/content/v1/607f89e638219e13eee71b1e/1684821560422-SD5V37BAG28BURTLIXUQ/michael-sum-LEpfefQf4rU-unsplash.jpg"
-            )
-            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="Colormap",
-                )
-            # Backbone selector (default = smaller/faster ViT-S/16+)
+        with gr.Column():
             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.",
+                info="ViT‑S/16+ is smaller & faster; ViT‑H/16+ is larger.",
             )
-            # Start processing button (manual trigger)
-            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)")
+            target_long_side = gr.Slider(224, 1024, value=768, step=16, label="Processing resolution")
+            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="Colormap")
+            exclude_r = gr.Slider(0, 10, value=0, step=1, label="Exclude radius (patches, only for Image 1)")
+            start_btn = gr.Button("▶️ Process both", variant="primary")
+
+        with gr.Column():
+            img1 = gr.Image(label="Image 1 (clickable after processing)", type="pil",
+                            value="https://images.squarespace-cdn.com/content/v1/607f89e638219e13eee71b1e/1684821560422-SD5V37BAG28BURTLIXUQ/michael-sum-LEpfefQf4rU-unsplash.jpg")
+            img2 = gr.Image(label="Image 2", type="pil",
+                            value="https://upload.wikimedia.org/wikipedia/commons/9/99/Golden_retriever_eating_pigs_foot.jpg")
 
     with gr.Row():
-        marked_image = gr.Image(label="Click marker / Preview", 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)
-
+        with gr.Column():
+            marked1 = gr.Image(label="Image 1 — click marker / preview", interactive=False)
+            heat1 = gr.Image(label="Image 1 — similarity heatmap", interactive=False)
+            overlay1 = gr.Image(label="Image 1 — overlay", interactive=False)
+        with gr.Column():
+            heat2 = gr.Image(label="Image 2 — similarity heatmap", interactive=False)
+            overlay2 = gr.Image(label="Image 2 — overlay", interactive=False)
+            score2 = gr.Number(label="Image 2 — max similarity score", precision=6)
+
+    # Utilities
     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
 
-    # Manual feature extraction (only runs on Start button)
-    def _run_extraction(img: Image.Image, long_side: int, model_id: str, progress=gr.Progress(track_tqdm=True)):
-        if img is None:
-            return None, None
+    # Process button: extract features for both images
+    def _run_both(im1: Image.Image, im2: Image.Image, long_side: int, model_id: str, progress=gr.Progress(track_tqdm=False)):
+        if im1 is None or im2 is None:
+            raise gr.Error("Please provide both images.")
         _ensure_model(model_id)
-        progress(0, desc="Extracting features...")
-        st = extract_image_features(img, int(long_side))
-        progress(1, desc="Done!")
-        return st["img"], st
-
-    # Clicking on processed image to compute similarities
-    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 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),
-        )
-
-    # On image change: just preview and clear outputs/state (NO extraction)
-    def _on_image_changed(img: Image.Image):
-        if img is None:
-            return None, None, None, None, None
-        return img, None, None, None, None
+        progress(0, desc="Extracting features…")
+        st1 = extract_image_features(im1, int(long_side), mdl=model)
+        st2 = extract_image_features(im2, int(long_side), mdl=model)
+        progress(1, desc="Done")
+        return st1["img"], st2["img"], st1, st2
 
-    # ---------- Wiring (Manual mode) ----------
-    # Do NOT auto-run on upload/slider/model change or on app load.
-    # Only the Start button triggers extraction.
     start_btn.click(
-        _run_extraction,
-        inputs=[input_image, target_long_side, model_choice],
-        outputs=[marked_image, app_state],
+        _run_both,
+        inputs=[img1, img2, target_long_side, model_choice],
+        outputs=[marked1, overlay2, state1, state2],  # show previews in two spots to confirm processing
     )
 
-    # When a new image is picked, show it as preview and clear old results.
-    input_image.change(
-        _on_image_changed,
-        inputs=[input_image],
-        outputs=[marked_image, app_state, heatmap_output, overlay_output, overlay_boxes_output],
-    )
+    # Clicking on Image 1
+    def _on_click(st1, st2, a: float, m: str, excl: int, evt: gr.SelectData):
+        if not st1 or not st2 or evt is None:
+            return (None,)*6
+        return click_two_image_similarity(
+            st1, st2, click_xy=evt.index,
+            exclude_radius_patches=int(excl),
+            alpha=float(a), cmap_name=m,
+        )
 
-    # Keep click handler the same.
-    marked_image.select(
+    marked1.select(
         _on_click,
-        inputs=[app_state, alpha, cmap, exclude_r, topk, box_radius],
-        outputs=[marked_image, heatmap_output, overlay_output, overlay_boxes_output],
+        inputs=[state1, state2, alpha, cmap, exclude_r],
+        outputs=[marked1, heat1, overlay1, heat2, overlay2, score2],
     )
 
 if __name__ == "__main__":