Create app.py

app.py

@@ -0,0 +1,276 @@
+# app.py
+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
+from matplotlib import colormaps
+from transformers import AutoModel # 💡 Import AutoModel
+
+# ----------------------------
+# Configuration
+# ----------------------------
+# 💡 Use the full, correct model ID from the Hugging Face Hub
+MODEL_ID = "facebook/dinov3-vith16plus-pretrain-lvd1689m"
+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)
+# ----------------------------
+def load_model_from_hub():
+    """Loads the DINOv3 model from the Hugging Face Hub."""
+    print(f"Loading model '{MODEL_ID}' from Hugging Face Hub...")
+    try:
+        # This will use the HF_TOKEN secret if you set it in your Space settings.
+        token = os.environ.get("HF_TOKEN")
+        # trust_remote_code is necessary for DINOv3
+        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
+    except Exception as e:
+        print(f"❌ Failed to load model: {e}")
+        # This will display a clear error message in the Gradio interface
+        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"Original error: {e}"
+        )
+
+# Load the model globally when the app starts
+model = load_model_from_hub()
+
+# ----------------------------
+# 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)))
+    new_w = max(patch, int(round(w * scale)))
+    new_h = ((new_h + patch - 1) // patch) * patch
+    new_w = ((new_w + patch - 1) // patch) * patch
+    return TF.to_tensor(TF.resize(img.convert("RGB"), (new_h, new_w)))
+
+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)
+    return Image.fromarray(rgb)
+
+def blend(base: Image.Image, heat: Image.Image, alpha: float = 0.55) -> Image.Image:
+    base = base.convert("RGBA")
+    heat = heat.convert("RGBA")
+    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)
+    out = img.copy()
+    draw = ImageDraw.Draw(out)
+    draw.line([(x - r, y), (x + r, y)], fill="red", width=3)
+    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)
+# ----------------------------
+@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
+    
+    # 💡 Use the standard forward pass of the transformers model
+    outputs = model(t_norm)
+    
+    # 💡 The model output includes a [CLS] token AND 4 register tokens.
+    # We must skip all 5 to get only the patch embeddings.
+    n_special_tokens = 5 
+    patch_embeddings = outputs.last_hidden_state.squeeze(0)[n_special_tokens:, :]
+    
+    # L2-normalize the features to prepare 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 inside the same image (No changes needed here)
+# ----------------------------
+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
+    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),
+            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"))
+
+    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)
+        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
+
+# ----------------------------
+# Gradio UI (No changes needed here)
+# ----------------------------
+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, then **click anywhere** to highlight the most similar regions in the *same* image.")
+
+    app_state = gr.State()
+
+    with gr.Row():
+        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"
+            )
+            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",
+                )
+        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=4, 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 _on_upload_or_slider_change(img: Image.Image, long_side: int, progress=gr.Progress(track_tqdm=True)):
+        if img is None:
+            return None, None
+        progress(0, desc="Extracting features...")
+        st = extract_image_features(img, int(long_side))
+        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 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),
+        )
+
+    # Wire events
+    inputs_for_update = [input_image, target_long_side]
+    outputs_for_update = [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
+
+    marked_image.select(
+        _on_click,
+        inputs=[app_state, alpha, cmap, exclude_r, topk, box_radius],
+        outputs=[marked_image, heatmap_output, overlay_output, overlay_boxes_output],
+    )
+
+if __name__ == "__main__":
+    demo.launch()