app.py

# 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

# --- 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
# ----------------------------
# 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) with caching
# ----------------------------
_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")  # 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}"
        )

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)))
    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 = (_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")
    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

    # 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 inside the same image
# ----------------------------
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 (Manual-only processing)
# ----------------------------
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.")

    app_state = 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+)
            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 (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)")

    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)

    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
        _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

    # ---------- 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],
    )

    # 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],
    )

    # Keep click handler 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()