{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# DINOv2 Tutorial — Notebook 11: Image Similarity & Retrieval\n", "\n", "**Series:** DINOv2 Deep Dive | **Prev:** `10_attention_head_specialization.ipynb` | **Next:** `12_downstream_tasks.ipynb`\n", "\n", "---\n", "\n", "## What You Will Learn\n", "- Computing cosine similarity between images using DINOv2 CLS features\n", "- Building a mini image retrieval system\n", "- Nearest-neighbor retrieval demo\n", "- Pairwise similarity matrix visualization\n", "- Ranking images by similarity to a query\n", "- How DINOv2 features enable zero-shot retrieval\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib.gridspec as gridspec\n", "from PIL import Image, ImageDraw\n", "from transformers import AutoImageProcessor, AutoModel\n", "\n", "MODEL_PATH = \"/mnt/mydata/huggingface/dinov2-base\"\n", "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "\n", "processor = AutoImageProcessor.from_pretrained(MODEL_PATH)\n", "model = AutoModel.from_pretrained(MODEL_PATH).to(DEVICE).eval()\n", "\n", "def extract_cls(image):\n", " inp = processor(images=image, return_tensors='pt')\n", " inp = {k: v.to(DEVICE) for k, v in inp.items()}\n", " with torch.no_grad():\n", " out = model(**inp)\n", " return out.last_hidden_state[0, 0].cpu().numpy()\n", "\n", "def cosine_similarity(a, b):\n", " a_n = a / (np.linalg.norm(a) + 1e-8)\n", " b_n = b / (np.linalg.norm(b) + 1e-8)\n", " return float(np.dot(a_n, b_n))\n", "\n", "print(\"Model loaded.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11.1 Build an Image Gallery" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Create a diverse image gallery ────────────────────────────────────────────\n", "def make_img(bg, shape, sc, size=224):\n", " img = Image.new('RGB', (size, size), bg)\n", " draw = ImageDraw.Draw(img)\n", " if shape == 'circle':\n", " draw.ellipse([30, 30, 194, 194], fill=sc)\n", " elif shape == 'square':\n", " draw.rectangle([30, 30, 194, 194], fill=sc)\n", " elif shape == 'triangle':\n", " draw.polygon([(112, 20), (20, 204), (204, 204)], fill=sc)\n", " elif shape == 'diamond':\n", " draw.polygon([(112, 20), (204, 112), (112, 204), (20, 112)], fill=sc)\n", " elif shape == 'cross':\n", " draw.rectangle([90, 20, 134, 204], fill=sc)\n", " draw.rectangle([20, 90, 204, 134], fill=sc)\n", " return img\n", "\n", "gallery = [\n", " # Red shapes\n", " (make_img((240,240,240), 'circle', (200,50,50)), 'Red Circle 1'),\n", " (make_img((220,220,220), 'circle', (180,40,40)), 'Red Circle 2'),\n", " (make_img((240,240,240), 'circle', (220,60,60)), 'Red Circle 3'),\n", " # Blue shapes\n", " (make_img((240,240,240), 'square', (50,80,200)), 'Blue Square 1'),\n", " (make_img((220,220,220), 'square', (40,70,180)), 'Blue Square 2'),\n", " (make_img((240,240,240), 'square', (60,90,220)), 'Blue Square 3'),\n", " # Green shapes\n", " (make_img((240,240,240), 'triangle', (50,180,50)), 'Green Triangle 1'),\n", " (make_img((220,220,220), 'triangle', (40,160,40)), 'Green Triangle 2'),\n", " (make_img((240,240,240), 'triangle', (60,200,60)), 'Green Triangle 3'),\n", " # Yellow shapes\n", " (make_img((240,240,240), 'diamond', (220,200,50)), 'Yellow Diamond 1'),\n", " (make_img((220,220,220), 'diamond', (200,180,40)), 'Yellow Diamond 2'),\n", " (make_img((240,240,240), 'cross', (200,180,50)), 'Yellow Cross'),\n", " # Mixed\n", " (make_img((200,230,200), 'circle', (50,150,50)), 'Green Circle'),\n", " (make_img((200,200,230), 'square', (50,50,180)), 'Blue Square Dark'),\n", " (make_img((230,200,200), 'triangle', (180,50,50)), 'Red Triangle'),\n", "]\n", "\n", "gallery_images = [img for img, _ in gallery]\n", "gallery_labels = [lbl for _, lbl in gallery]\n", "n_gallery = len(gallery)\n", "\n", "# Extract features\n", "print(\"Extracting features...\")\n", "gallery_features = np.array([extract_cls(img) for img in gallery_images])\n", "print(f\"Gallery features: {gallery_features.shape}\")\n", "\n", "# Normalize\n", "gallery_norms = np.linalg.norm(gallery_features, axis=1, keepdims=True)\n", "gallery_features_norm = gallery_features / (gallery_norms + 1e-8)\n", "\n", "# Show gallery\n", "fig, axes = plt.subplots(3, 5, figsize=(18, 11))\n", "for i, (img, label) in enumerate(gallery):\n", " axes[i//5, i%5].imshow(img)\n", " axes[i//5, i%5].set_title(f'[{i}] {label}', fontsize=8, fontweight='bold')\n", " axes[i//5, i%5].axis('off')\n", "plt.suptitle('Image Gallery (15 images)', fontsize=13, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11.2 Pairwise Similarity Matrix" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Full pairwise similarity matrix ───────────────────────────────────────────\n", "sim_matrix = gallery_features_norm @ gallery_features_norm.T # [15, 15]\n", "\n", "fig, axes = plt.subplots(1, 2, figsize=(18, 8))\n", "\n", "im = axes[0].imshow(sim_matrix, cmap='RdYlGn', vmin=0.5, vmax=1.0)\n", "axes[0].set_xticks(range(n_gallery))\n", "axes[0].set_xticklabels([f'[{i}]' for i in range(n_gallery)], rotation=45, ha='right')\n", "axes[0].set_yticks(range(n_gallery))\n", "axes[0].set_yticklabels([f'[{i}] {gallery_labels[i][:12]}' for i in range(n_gallery)], fontsize=8)\n", "for i in range(n_gallery):\n", " for j in range(n_gallery):\n", " axes[0].text(j, i, f'{sim_matrix[i,j]:.2f}',\n", " ha='center', va='center', fontsize=6)\n", "plt.colorbar(im, ax=axes[0])\n", "axes[0].set_title('Pairwise Cosine Similarity Matrix\\n(DINOv2 CLS features)',\n", " fontweight='bold')\n", "\n", "# Distribution of similarities\n", "upper_tri = sim_matrix[np.triu_indices(n_gallery, k=1)]\n", "axes[1].hist(upper_tri, bins=20, color='steelblue', alpha=0.8, edgecolor='black')\n", "axes[1].axvline(upper_tri.mean(), color='red', linestyle='--', linewidth=2,\n", " label=f'Mean = {upper_tri.mean():.3f}')\n", "axes[1].axvline(upper_tri.median() if hasattr(upper_tri, 'median') else np.median(upper_tri),\n", " color='orange', linestyle='--', linewidth=2,\n", " label=f'Median = {np.median(upper_tri):.3f}')\n", "axes[1].set_title('Distribution of Pairwise Similarities\\n(upper triangle)', fontweight='bold')\n", "axes[1].set_xlabel('Cosine Similarity')\n", "axes[1].set_ylabel('Count')\n", "axes[1].legend()\n", "axes[1].grid(True, alpha=0.3)\n", "\n", "plt.suptitle('Image Gallery: Pairwise Similarity Analysis', fontsize=13, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "# Most and least similar pairs\n", "np.fill_diagonal(sim_matrix, -1) # exclude self\n", "most_sim_idx = np.unravel_index(sim_matrix.argmax(), sim_matrix.shape)\n", "least_sim_idx = np.unravel_index(sim_matrix.argmin(), sim_matrix.shape)\n", "np.fill_diagonal(sim_matrix, 1.0) # restore\n", "\n", "print(f\"Most similar pair: [{most_sim_idx[0]}] {gallery_labels[most_sim_idx[0]]} \"\n", " f\"<-> [{most_sim_idx[1]}] {gallery_labels[most_sim_idx[1]]} \"\n", " f\"(sim={sim_matrix[most_sim_idx]:.4f})\")\n", "print(f\"Least similar pair: [{least_sim_idx[0]}] {gallery_labels[least_sim_idx[0]]} \"\n", " f\"<-> [{least_sim_idx[1]}] {gallery_labels[least_sim_idx[1]]} \"\n", " f\"(sim={sim_matrix[least_sim_idx]:.4f})\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11.3 Image Retrieval: Query → Top-K Results" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Retrieval function ────────────────────────────────────────────────────────\n", "def retrieve_top_k(query_feat, gallery_feats_norm, gallery_labels, k=5, exclude_self=True):\n", " q_norm = query_feat / (np.linalg.norm(query_feat) + 1e-8)\n", " sims = gallery_feats_norm @ q_norm # [N]\n", " if exclude_self:\n", " sims[sims >= 0.9999] = -1 # exclude exact match (self)\n", " top_k_idx = sims.argsort()[::-1][:k]\n", " return top_k_idx, sims[top_k_idx]\n", "\n", "# Run retrieval for 3 query images\n", "query_indices = [0, 3, 6] # Red Circle 1, Blue Square 1, Green Triangle 1\n", "k = 5\n", "\n", "fig, axes = plt.subplots(len(query_indices), k+1, figsize=(22, 12))\n", "\n", "for row, q_idx in enumerate(query_indices):\n", " query_feat = gallery_features[q_idx]\n", " top_k_idx, top_k_sims = retrieve_top_k(\n", " query_feat, gallery_features_norm, gallery_labels, k=k\n", " )\n", " \n", " # Query image\n", " axes[row, 0].imshow(gallery_images[q_idx])\n", " axes[row, 0].set_title(f'QUERY\\n[{q_idx}] {gallery_labels[q_idx]}',\n", " fontweight='bold', color='blue', fontsize=9)\n", " axes[row, 0].axis('off')\n", " # Red border for query\n", " for spine in axes[row, 0].spines.values():\n", " spine.set_edgecolor('blue'); spine.set_linewidth(4)\n", " \n", " # Retrieved images\n", " for col, (ret_idx, sim) in enumerate(zip(top_k_idx, top_k_sims)):\n", " axes[row, col+1].imshow(gallery_images[ret_idx])\n", " axes[row, col+1].set_title(\n", " f'Rank {col+1}\\n[{ret_idx}] {gallery_labels[ret_idx][:12]}\\nsim={sim:.3f}',\n", " fontsize=8, fontweight='bold'\n", " )\n", " axes[row, col+1].axis('off')\n", " # Color border by similarity\n", " border_color = 'green' if sim > 0.9 else ('orange' if sim > 0.8 else 'red')\n", " for spine in axes[row, col+1].spines.values():\n", " spine.set_edgecolor(border_color); spine.set_linewidth(3)\n", "\n", "plt.suptitle(f'Image Retrieval: Top-{k} Results per Query\\n'\n", " 'Green border = high similarity (>0.9), Orange = medium, Red = low',\n", " fontsize=13, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11.4 Similarity Ranking Visualization" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Rank all gallery images by similarity to a query ──────────────────────────\n", "query_idx = 0 # Red Circle 1\n", "query_feat = gallery_features[query_idx]\n", "q_norm = query_feat / (np.linalg.norm(query_feat) + 1e-8)\n", "all_sims = gallery_features_norm @ q_norm # [15]\n", "\n", "# Sort by similarity\n", "sorted_idx = all_sims.argsort()[::-1]\n", "\n", "fig = plt.figure(figsize=(20, 10))\n", "gs = gridspec.GridSpec(2, 8, figure=fig)\n", "\n", "# Query image (large)\n", "ax_query = fig.add_subplot(gs[:, 0])\n", "ax_query.imshow(gallery_images[query_idx])\n", "ax_query.set_title(f'QUERY\\n[{query_idx}]\\n{gallery_labels[query_idx]}',\n", " fontweight='bold', color='blue', fontsize=10)\n", "ax_query.axis('off')\n", "\n", "# Ranked results\n", "for rank, idx in enumerate(sorted_idx):\n", " if rank >= 14:\n", " break\n", " row = rank // 7\n", " col = rank % 7 + 1\n", " ax = fig.add_subplot(gs[row, col])\n", " ax.imshow(gallery_images[idx])\n", " sim = all_sims[idx]\n", " is_self = (idx == query_idx)\n", " color = 'gold' if is_self else ('green' if sim > 0.9 else ('orange' if sim > 0.8 else 'red'))\n", " ax.set_title(f'#{rank+1} [{idx}]\\n{gallery_labels[idx][:10]}\\n{sim:.3f}',\n", " fontsize=7, fontweight='bold', color=color)\n", " ax.axis('off')\n", " for spine in ax.spines.values():\n", " spine.set_edgecolor(color); spine.set_linewidth(2)\n", "\n", "plt.suptitle(f'All Gallery Images Ranked by Similarity to Query [{query_idx}]: {gallery_labels[query_idx]}',\n", " fontsize=13, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "# Print ranking\n", "print(f\"Similarity ranking for query [{query_idx}] {gallery_labels[query_idx]}:\")\n", "print(f\"{'Rank':>5} {'Idx':>5} {'Label':>20} {'Similarity':>12}\")\n", "print(\"-\" * 50)\n", "for rank, idx in enumerate(sorted_idx):\n", " marker = ' <-- SELF' if idx == query_idx else ''\n", " print(f\"{rank+1:>5} {idx:>5} {gallery_labels[idx]:>20} {all_sims[idx]:>12.4f}{marker}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11.5 Retrieval with a New Query Image" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Query with a new image not in the gallery ──────────────────────────────────\n", "# Create a query image similar to red circles but slightly different\n", "query_new = make_img((230, 230, 230), 'circle', (190, 45, 45))\n", "query_new_feat = extract_cls(query_new)\n", "q_new_norm = query_new_feat / (np.linalg.norm(query_new_feat) + 1e-8)\n", "new_sims = gallery_features_norm @ q_new_norm\n", "top5_idx = new_sims.argsort()[::-1][:5]\n", "\n", "fig, axes = plt.subplots(1, 6, figsize=(20, 4))\n", "\n", "axes[0].imshow(query_new)\n", "axes[0].set_title('NEW QUERY\\n(not in gallery)', fontweight='bold', color='blue')\n", "axes[0].axis('off')\n", "\n", "for col, idx in enumerate(top5_idx):\n", " axes[col+1].imshow(gallery_images[idx])\n", " axes[col+1].set_title(f'Rank {col+1}\\n{gallery_labels[idx]}\\nsim={new_sims[idx]:.3f}',\n", " fontsize=9, fontweight='bold')\n", " axes[col+1].axis('off')\n", "\n", "plt.suptitle('Zero-Shot Retrieval: New Query Image vs Gallery',\n", " fontsize=13, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11.6 Similarity Threshold Analysis" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Analyze retrieval precision at different thresholds ───────────────────────\n", "# Define ground truth: images with same shape type are relevant\n", "shape_groups = {\n", " 'circle': [0, 1, 2, 12],\n", " 'square': [3, 4, 5, 13],\n", " 'triangle': [6, 7, 8, 14],\n", " 'diamond': [9, 10],\n", " 'cross': [11],\n", "}\n", "\n", "def get_shape_group(idx):\n", " for shape, indices in shape_groups.items():\n", " if idx in indices:\n", " return shape\n", " return 'unknown'\n", "\n", "thresholds = np.linspace(0.5, 1.0, 50)\n", "precisions = []\n", "recalls = []\n", "\n", "for threshold in thresholds:\n", " tp = fp = fn = 0\n", " for q_idx in range(n_gallery):\n", " q_shape = get_shape_group(q_idx)\n", " q_sims = sim_matrix[q_idx].copy()\n", " q_sims[q_idx] = 0 # exclude self\n", " retrieved = np.where(q_sims >= threshold)[0]\n", " relevant = [i for i in range(n_gallery)\n", " if i != q_idx and get_shape_group(i) == q_shape]\n", " tp += len([i for i in retrieved if i in relevant])\n", " fp += len([i for i in retrieved if i not in relevant])\n", " fn += len([i for i in relevant if i not in retrieved])\n", " precision = tp / (tp + fp + 1e-8)\n", " recall = tp / (tp + fn + 1e-8)\n", " precisions.append(precision)\n", " recalls.append(recall)\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(18, 6))\n", "\n", "axes[0].plot(thresholds, precisions, 'b-', linewidth=2, label='Precision')\n", "axes[0].plot(thresholds, recalls, 'r-', linewidth=2, label='Recall')\n", "axes[0].set_title('Precision & Recall vs Similarity Threshold', fontweight='bold')\n", "axes[0].set_xlabel('Similarity Threshold')\n", "axes[0].set_ylabel('Score')\n", "axes[0].legend(); axes[0].grid(True, alpha=0.3)\n", "\n", "# F1 score\n", "f1_scores = [2*p*r/(p+r+1e-8) for p, r in zip(precisions, recalls)]\n", "best_threshold = thresholds[np.argmax(f1_scores)]\n", "axes[1].plot(thresholds, f1_scores, 'g-', linewidth=2)\n", "axes[1].axvline(best_threshold, color='red', linestyle='--',\n", " label=f'Best threshold = {best_threshold:.3f}')\n", "axes[1].set_title('F1 Score vs Similarity Threshold', fontweight='bold')\n", "axes[1].set_xlabel('Similarity Threshold')\n", "axes[1].set_ylabel('F1 Score')\n", "axes[1].legend(); axes[1].grid(True, alpha=0.3)\n", "\n", "# Precision-Recall curve\n", "axes[2].plot(recalls, precisions, 'purple', linewidth=2)\n", "axes[2].fill_between(recalls, precisions, alpha=0.2, color='purple')\n", "axes[2].set_title('Precision-Recall Curve', fontweight='bold')\n", "axes[2].set_xlabel('Recall')\n", "axes[2].set_ylabel('Precision')\n", "axes[2].set_xlim(0, 1); axes[2].set_ylim(0, 1)\n", "axes[2].grid(True, alpha=0.3)\n", "auc = np.trapz(precisions[::-1], recalls[::-1])\n", "axes[2].text(0.5, 0.1, f'AUC = {auc:.3f}', ha='center', fontsize=12,\n", " transform=axes[2].transAxes, fontweight='bold')\n", "\n", "plt.suptitle('Retrieval Performance Analysis\\n'\n", " '(ground truth: same shape type = relevant)',\n", " fontsize=13, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "print(f\"Best similarity threshold: {best_threshold:.3f}\")\n", "print(f\"Best F1 score: {max(f1_scores):.3f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11.7 Summary\n", "\n", "| Task | Method | Key Metric |\n", "|------|--------|------------|\n", "| Image similarity | Cosine similarity of CLS features | Score in [0,1] |\n", "| Image retrieval | k-NN search in feature space | Precision@k, Recall@k |\n", "| Threshold selection | F1 score optimization | Best threshold |\n", "| Zero-shot retrieval | No training needed | Works out of the box |\n", "\n", "**Key insight:** DINOv2 CLS features enable **zero-shot image retrieval** — no fine-tuning needed. \n", "The features are semantically meaningful enough to find similar images by content.\n", "\n", "---\n", "**Next:** [`12_downstream_tasks.ipynb`](12_downstream_tasks.ipynb)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 5 }