Spaces:
Running
Running
| 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) | |
| # ---------------------------- | |
| 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 | |
| # ---------------------------- | |
| 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("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="1. Upload Image", | |
| 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", | |
| ) | |
| # ⭐️ 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( | |
| ["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=1, step=1, label="Box radius (patches)") | |
| with gr.Row(): | |
| 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) | |
| # ⭐️ 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! 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), | |
| topk=int(k), alpha=float(a), cmap_name=m, | |
| box_radius_patches=int(box_rad), | |
| ) | |
| # ⭐️ UPDATED EVENT WIRING | |
| inputs_for_processing = [input_image, target_long_side] | |
| outputs_for_processing = [marked_image, app_state] | |
| # 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], | |
| outputs=[marked_image, heatmap_output, overlay_output, overlay_boxes_output], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |