{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Visualizing ViT Patch Features with PCA\n",
    "\n",
    "[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/pca_feature_visualization.ipynb)\n",
    "\n",
    "Vision Transformers (ViTs) split an image into a grid of patches and produce one feature\n",
    "vector per patch. These *patch features* (or *dense features*) capture what the model\n",
    "understands about each region of the image. But each vector has hundreds of dimensions,\n",
    "so we can't look at them directly.\n",
    "\n",
    "A simple and surprisingly effective trick, popularized by\n",
    "[DINO](https://arxiv.org/abs/2104.14294), is to run **PCA** on the patch features and map\n",
    "the top three principal components to the red, green, and blue channels of an image. The\n",
    "result paints each patch by its dominant features: semantically similar regions get\n",
    "similar colors, so objects pop out from the background without the model ever being\n",
    "trained for segmentation. Because of this, PCA maps have become a standard qualitative\n",
    "check of feature quality for self-supervised and distilled backbones, and the DINOv2 and\n",
    "DINOv3 papers both use them to showcase their models.\n",
    "\n",
    "In this tutorial we:\n",
    "\n",
    "1. Load a pretrained **DINOv3** backbone with LightlyTrain and extract its patch features.\n",
    "2. Turn those features into an RGB image with PCA.\n",
    "3. Compare the feature maps of **DINOv3**, **DINOv2**, and **EUPE** on the same image.\n",
    "4. Meet Lightly's distilled **ViT-Tiny** models and see how well their features hold up\n",
    "   against much larger backbones."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Install Dependencies\n",
    "\n",
    "We only need `lightly-train`. It already pulls in `torch`, `numpy`, and `matplotlib`,\n",
    "which is everything we use here. The PCA is computed directly with `torch`, so there is\n",
    "no need for scikit-learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install lightly-train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Get a Sample Image\n",
    "\n",
    "Let's download a single image to visualize. Feel free to swap in your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import urllib.request\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "url = \"https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg\"\n",
    "image = Image.open(io.BytesIO(urllib.request.urlopen(url).read())).convert(\"RGB\")\n",
    "image.resize((512, 512))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Preprocess the Image\n",
    "\n",
    "Before feeding the image to the model we resize it to a multiple of the model's patch\n",
    "size and normalize it with the standard ImageNet statistics that the DINO models were\n",
    "trained with. Resizing to a multiple of the patch size means the image divides evenly\n",
    "into patches. A larger size gives a finer feature grid (and a sharper visualization) at\n",
    "the cost of more compute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.transforms.v2 as transforms\n",
    "\n",
    "# Standard ImageNet normalization used by the DINOv2/DINOv3 backbones.\n",
    "IMAGENET_MEAN = (0.485, 0.456, 0.406)\n",
    "IMAGENET_STD = (0.229, 0.224, 0.225)\n",
    "\n",
    "\n",
    "def preprocess(image, patch_size, size=768):\n",
    "    # Round the target size down to a multiple of the patch size.\n",
    "    size = (size // patch_size) * patch_size\n",
    "    transform = transforms.Compose(\n",
    "        [\n",
    "            transforms.Resize((size, size)),\n",
    "            transforms.ToImage(),\n",
    "            transforms.ToDtype(torch.float32, scale=True),\n",
    "            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),\n",
    "        ]\n",
    "    )\n",
    "    return transform(image).unsqueeze(0)  # (1, 3, H, W)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Load a DINOv3 Backbone and Extract Patch Features\n",
    "\n",
    "We load a pretrained DINOv3 backbone with LightlyTrain. `get_wrapped_model` accepts any\n",
    "model name from `lightly_train.list_models()` (e.g. `\"dinov3/vits16\"`, `\"dinov2/vitb14\"`)\n",
    "and downloads the pretrained weights on first use. Calling `forward_features` on the\n",
    "returned wrapper gives us the patch features as a spatial grid of shape\n",
    "`(1, feature_dim, H, W)`, where `H` and `W` are the number of patches along each axis.\n",
    "\n",
    "```{note}\n",
    "`get_wrapped_model` lives under `lightly_train._models` and is an internal API that may\n",
    "change between releases. The public `embed` command only returns a single pooled vector\n",
    "per image, so we use the wrapper directly here to access the per-patch features.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lightly_train._models.package_helpers import get_wrapped_model\n",
    "\n",
    "# Load a pretrained DINOv3 backbone (weights download on first use).\n",
    "wrapper = get_wrapped_model(\"dinov3/vitb16\", num_input_channels=3).eval()\n",
    "\n",
    "x = preprocess(image, patch_size=wrapper.patch_size())\n",
    "with torch.no_grad():\n",
    "    features = wrapper.forward_features(x)[\"features\"]  # (1, feature_dim, H, W)\n",
    "\n",
    "print(f\"patch size:   {wrapper.patch_size()}\")\n",
    "print(f\"feature dim:  {wrapper.feature_dim()}\")\n",
    "print(f\"feature grid: {tuple(features.shape)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Reduce the Features to RGB with PCA\n",
    "\n",
    "Now the core step. We treat every patch as a data point in `feature_dim`-dimensional\n",
    "space, center the points, and project them onto their top three principal components with\n",
    "`torch.pca_lowrank`. Those three components become the R, G, and B channels. Finally we\n",
    "normalize each channel to the 1st to 99th percentile range so a few outlier patches don't\n",
    "wash out the colors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "def pca_rgb(features):\n",
    "    \"\"\"Maps ViT patch features to an (H, W, 3) RGB image via PCA.\n",
    "\n",
    "    Args:\n",
    "        features: Patch features of shape (1, feature_dim, H, W).\n",
    "\n",
    "    Returns:\n",
    "        An (H, W, 3) float array with values in [0, 1].\n",
    "    \"\"\"\n",
    "    _, feature_dim, h, w = features.shape\n",
    "    # Flatten the spatial grid into a (num_patches, feature_dim) matrix.\n",
    "    tokens = features[0].permute(1, 2, 0).reshape(-1, feature_dim)\n",
    "    tokens = tokens - tokens.mean(dim=0, keepdim=True)\n",
    "\n",
    "    # Project onto the top 3 principal components.\n",
    "    _, _, v = torch.pca_lowrank(tokens, q=3)\n",
    "    projected = (tokens @ v[:, :3]).reshape(h, w, 3).cpu().numpy()\n",
    "\n",
    "    # Percentile-normalize each channel to [0, 1] for visualization.\n",
    "    low = np.percentile(projected, 1, axis=(0, 1), keepdims=True)\n",
    "    high = np.percentile(projected, 99, axis=(0, 1), keepdims=True)\n",
    "    return np.clip((projected - low) / (high - low + 1e-8), 0, 1)\n",
    "\n",
    "\n",
    "rgb = pca_rgb(features)\n",
    "print(f\"PCA image shape: {rgb.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Show the Result\n",
    "\n",
    "The PCA image is small (one pixel per patch), so we upsample it to the input resolution\n",
    "with bilinear interpolation and plot it next to the original image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "def upsample_to(rgb, size):\n",
    "    # rgb: (H, W, 3) -> upsample to (size, size, 3).\n",
    "    grid = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)  # (1, 3, H, W)\n",
    "    grid = F.interpolate(grid, size=size, mode=\"bilinear\", align_corners=False)\n",
    "    return grid[0].permute(1, 2, 0).numpy()\n",
    "\n",
    "\n",
    "size = x.shape[-1]\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "axes[0].imshow(image.resize((size, size)))\n",
    "axes[0].set_title(\"Input\")\n",
    "axes[1].imshow(upsample_to(rgb, size))\n",
    "axes[1].set_title(\"DINOv3 patch features (PCA)\")\n",
    "for ax in axes:\n",
    "    ax.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Why PCA Maps Are a Useful Quality Check\n",
    "\n",
    "SSL backbones are trained without labels, so there is no validation accuracy to glance\n",
    "at, and proper evaluations like linear probing, k-NN, or fine-tuning on a downstream\n",
    "task require labeled data and GPU hours. PCA maps fill this gap as a qualitative metric,\n",
    "for three reasons:\n",
    "\n",
    "1. **They probe exactly what dense downstream tasks consume.** Detection, segmentation,\n",
    "   and depth heads are trained on top of the patch features. If objects already emerge\n",
    "   as coherent, sharply bounded regions without any task training, a small head has an\n",
    "   easy job; if the features are noisy or entangled with the background, the head must\n",
    "   compensate.\n",
    "2. **They are label-free and instant.** You can run them in seconds on images from your\n",
    "   own domain, making them ideal as a first sanity check of a pretrained or freshly\n",
    "   distilled checkpoint before committing to a full evaluation.\n",
    "3. **They expose failure modes that global benchmarks hide.** A backbone can score well\n",
    "   on classification while its dense features are degraded: high-norm\n",
    "   \"outlier\" patches show up as speckle noise, and\n",
    "   weak localization shows up as blurry object boundaries.\n",
    "\n",
    "Keep in mind that the maps are qualitative: they complement, but do not replace, a\n",
    "quantitative evaluation on your downstream task."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Compare Backbones\n",
    "\n",
    "The exact same pipeline works for any ViT backbone that LightlyTrain ships. Let's run it\n",
    "for DINOv3, DINOv2, and [EUPE](https://huggingface.co/facebook/EUPE-ViT-S) side by side.\n",
    "Each model computes its own PCA basis, so the colors are not directly comparable between\n",
    "panels, what to look for is how cleanly each model separates the object from the\n",
    "background and how much fine detail its features capture.\n",
    "\n",
    "Run `lightly_train.list_models()` to see every available backbone name you can drop into\n",
    "this loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {\n",
    "    \"DINOv3\": \"dinov3/vitb16\",\n",
    "    \"DINOv2\": \"dinov2/vitb14\",\n",
    "    \"EUPE\": \"dinov3/vitb16-eupe\",\n",
    "}\n",
    "\n",
    "fig, axes = plt.subplots(1, len(models) + 1, figsize=(5 * (len(models) + 1), 5))\n",
    "axes[0].imshow(image.resize((size, size)))\n",
    "axes[0].set_title(\"Input\")\n",
    "axes[0].axis(\"off\")\n",
    "\n",
    "for ax, (title, name) in zip(axes[1:], models.items()):\n",
    "    wrapper = get_wrapped_model(name, num_input_channels=3).eval()\n",
    "    x = preprocess(image, patch_size=wrapper.patch_size())\n",
    "    with torch.no_grad():\n",
    "        features = wrapper.forward_features(x)[\"features\"]\n",
    "    ax.imshow(upsample_to(pca_rgb(features), size))\n",
    "    ax.set_title(title)\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## Smaller Models: Lightly's Distilled ViT-Tiny\n",
    "\n",
    "Large backbones like DINOv3 ViT-L produce excellent features, but they are often too\n",
    "heavy for real-time or embedded applications. For those cases Lightly offers tiny DINOv3\n",
    "models, trained with LightlyTrain's\n",
    "[distillation method](https://docs.lightly.ai/train/stable/pretrain_distill/methods/distillation.html)\n",
    "using DINOv3 ViT-L/16 as the teacher on ImageNet-1K. They ship ready to use as\n",
    "`dinov3/vitt16` and `dinov3/vitt16plus`, see the\n",
    "[DINOv3 model docs](https://docs.lightly.ai/train/stable/pretrain_distill/models/dinov3.html)\n",
    "for details. You can also distill your own tiny model on your own dataset with the same\n",
    "method.\n",
    "\n",
    "PCA maps are particularly handy for judging distilled models: the goal of distillation is\n",
    "to transfer the teacher's representation into a smaller student, so a good student should\n",
    "reproduce the teacher's feature structure, and any degradation is immediately visible as\n",
    "noisier maps or washed-out object boundaries, long before you run a full benchmark.\n",
    "\n",
    "Let's run the exact same PCA pipeline for DINOv3 ViT-B and the two tiny models.\n",
    "Despite having only a fraction of the parameters, the tiny models preserve the feature\n",
    "structure of the much larger models remarkably well: the object still separates cleanly\n",
    "from the background, so their features are a strong basis for downstream dense tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {\n",
    "    \"DINOv3 ViT-B\": \"dinov3/vitb16\",\n",
    "    \"ViT-Tiny+\": \"dinov3/vitt16plus\",\n",
    "    \"ViT-Tiny\": \"dinov3/vitt16\",\n",
    "}\n",
    "\n",
    "fig, axes = plt.subplots(1, len(models) + 1, figsize=(5 * (len(models) + 1), 5))\n",
    "axes[0].imshow(image.resize((size, size)))\n",
    "axes[0].set_title(\"Input\")\n",
    "axes[0].axis(\"off\")\n",
    "\n",
    "for ax, (title, name) in zip(axes[1:], models.items()):\n",
    "    wrapper = get_wrapped_model(name, num_input_channels=3).eval()\n",
    "    num_params = sum(p.numel() for p in wrapper.get_model().parameters())\n",
    "    x = preprocess(image, patch_size=wrapper.patch_size())\n",
    "    with torch.no_grad():\n",
    "        features = wrapper.forward_features(x)[\"features\"]\n",
    "    ax.imshow(upsample_to(pca_rgb(features), size))\n",
    "    ax.set_title(f\"{title}: {num_params / 1e6:.0f}M params\")\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this tutorial we loaded pretrained ViT backbones with LightlyTrain, extracted their\n",
    "per-patch features, and visualized them by projecting the features onto their top three\n",
    "principal components as RGB. This is a quick, label-free way to inspect what a foundation\n",
    "model has learned about your images: objects and coherent regions emerge as blocks of\n",
    "consistent color, even though the models were never trained for segmentation.\n",
    "\n",
    "We also saw that the Lightly-distilled `dinov3/vitt16` and `dinov3/vitt16plus` models\n",
    "retain most of the feature quality of their much larger teacher, making them a great\n",
    "starting point when you need a small and fast backbone, see the\n",
    "[DINOv3 model docs](https://docs.lightly.ai/train/stable/pretrain_distill/models/dinov3.html)\n",
    "to get started.\n",
    "\n",
    "From here you can swap in any backbone from `lightly_train.list_models()`, try your own\n",
    "images, or fine-tune these backbones on a downstream task. See the\n",
    "[LightlyTrain documentation](https://docs.lightly.ai/train/stable/) for detection,\n",
    "segmentation, and classification workflows."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}