{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p63J-QmUrMN-"
   },
   "outputs": [],
   "source": [
    "# Copyright (c) TorchGeo Contributors. All rights reserved.\n",
    "# Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XRSkMFqyrMOE"
   },
   "source": [
    "# Pretrained Weights\n",
    "\n",
    "_Written by: Caleb Robinson_\n",
    "\n",
    "In this tutorial, we will demonstrate how to use pretrained models in TorchGeo to extract fixed-length embeddings from remote sensing imagery. This is useful for quickly exploring datasets, visualizing feature spaces, and establishing baseline performance without the need for extensive training.\n",
    "\n",
    "Specifically, we will:\n",
    "\n",
    "- Load the **EuroSAT** dataset using TorchGeo's Lightning `DataModule`.\n",
    "- Load two pretrained encoders:\n",
    "  1) **DOFA** (a ViT-B/16-style encoder fine-tuned for Earth observation) - outputs 768-D features.\n",
    "  2) **ResNet-18** (using weights from SSL4EO) - outputs 512-D features.\n",
    "- Extract fixed-length embeddings for every image without labels (no gradient / just forward passes).\n",
    "- Train a simple k-Nearest Neighbors classifier on the embeddings to quantify linear separability.\n",
    "- Visualize the feature space with a 2-D PCA plot to see class separation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NBa5RPAirMOF"
   },
   "source": [
    "## Setup\n",
    "\n",
    "First, we install TorchGeo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5AIQ1B9DrMOG",
    "outputId": "6bf360ea-8f60-45cf-c96e-0eac54818079"
   },
   "outputs": [],
   "source": [
    "# On Colab, this ensures the latest TorchGeo is available.\n",
    "\n",
    "%pip install torchgeo scikit-learn tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IcCOnzVLrMOI"
   },
   "source": [
    "## Imports\n",
    "\n",
    "Next, we import TorchGeo and any other libraries we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rjEGiiurrMOI"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "\n",
    "import kornia.augmentation as K\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.metrics import classification_report\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from tqdm import tqdm\n",
    "\n",
    "from torchgeo.datamodules import EuroSAT100DataModule\n",
    "from torchgeo.datasets import EuroSAT100\n",
    "from torchgeo.models import DOFABase16_Weights, ResNet18_Weights, get_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QNnDoIf2rMOK"
   },
   "source": [
    "## Datamodule\n",
    "\n",
    "We will utilize TorchGeo's [Lightning](https://lightning.ai/docs/pytorch/stable/) datamodules to organize the dataloader setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ia5ktOVerMOL"
   },
   "outputs": [],
   "source": [
    "# Build EuroSAT DataModule (train/val loaders) using TorchGeo Lightning utilities.\n",
    "\n",
    "root = os.path.join(tempfile.gettempdir(), 'eurosat100')\n",
    "# ensure that the dataset is download\n",
    "EuroSAT100(root=root, split='train', download=True)\n",
    "EuroSAT100(root=root, split='val', download=True)\n",
    "\n",
    "datamodule = EuroSAT100DataModule(\n",
    "    root=root, batch_size=10, num_workers=2, bands=('B02', 'B03', 'B04')\n",
    ")\n",
    "datamodule.setup('fit')\n",
    "datamodule.setup('validate')\n",
    "\n",
    "train_dl = datamodule.train_dataloader()\n",
    "val_dl = datamodule.val_dataloader()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vgNswWKOrMOO"
   },
   "source": [
    "## Embedding\n",
    "\n",
    "We will embed the entirety of EuroSAT train and validation splits using a pretrained model by extracting features from the final layer before the classification head.\n",
    "\n",
    "With the DOFA model this will give us 768-dimensional feature vectors from each image. We use these vectors with the labels to train and evaluate a simple k-nearest neighbors classifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0Sf-CBorrMOO"
   },
   "outputs": [],
   "source": [
    "accelerator = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the DOFA pretrained encoder (ViT-style) and move it to device.\n",
    "\n",
    "model = get_model('dofa_base_patch16_224', weights=DOFABase16_Weights.DOFA_MAE)\n",
    "model = model.eval().to(accelerator)\n",
    "\n",
    "augs = K.AugmentationSequential(\n",
    "    K.Normalize(\n",
    "        mean=0.0, std=10_000, p=1.0\n",
    "    ),  # Divide by 10,000 to approximately scale to [0, 1]\n",
    "    K.Resize((224, 224), antialias=True),  # DOFA model expects 224x224 inputs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def embed_dofa(model, dataloader, accelerator, transforms=None):\n",
    "    \"\"\"Helper function to embed the samples from a dataloader using a DOFA model.\"\"\"\n",
    "    x_all = []\n",
    "    y_all = []\n",
    "    for batch in tqdm(dataloader):\n",
    "        x = batch['image'].to(accelerator)\n",
    "        y = batch['label']\n",
    "        x = x.to(accelerator)\n",
    "        if transforms is not None:\n",
    "            x = transforms(x)\n",
    "\n",
    "        with torch.inference_mode():\n",
    "            # DOFA requires us to forward the central wavelengths of each bands\n",
    "            # these are B02, B03, B04 wavelengths for Sentinel-2\n",
    "            embeddings = model.forward_features(x, wavelengths=[0.49, 0.56, 0.66])\n",
    "        x_all.append(embeddings.cpu().numpy())\n",
    "        y_all.append(y.numpy())\n",
    "    x_all = np.concatenate(x_all, axis=0)\n",
    "    y_all = np.concatenate(y_all, axis=0)\n",
    "    return x_all, y_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train, y_train = embed_dofa(model, train_dl, accelerator, transforms=augs)\n",
    "x_val, y_val = embed_dofa(model, val_dl, accelerator, transforms=augs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit a k-NN classifier on DOFA train embeddings and evaluate on validation embeddings. This gives a quick, label-efficient baseline without fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "knn_model = KNeighborsClassifier(n_neighbors=5)\n",
    "knn_model.fit(x_train, y_train)\n",
    "y_pred = knn_model.predict(x_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = [\n",
    "    'Annual Crop',\n",
    "    'Forest',\n",
    "    'Herbaceous Vegetation',\n",
    "    'Highway',\n",
    "    'Industrial Buildings',\n",
    "    'Pasture',\n",
    "    'Permanent Crop',\n",
    "    'Residential Buildings',\n",
    "    'River',\n",
    "    'Sea & Lake',\n",
    "]\n",
    "\n",
    "print(\n",
    "    classification_report(\n",
    "        y_val, y_pred, digits=2, target_names=class_names, zero_division=0\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's do the same thing with a ResNet18 model pretrained on Sentinel-2 RGB imagery (from the SSL4EO paper).\n",
    "\n",
    "We reuse the same dataloaders and evaluation code to isolate the effect of the encoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can instantiate a ResNet18 model in a similar way to the DOFA model above.\n",
    "\n",
    "model = get_model('resnet18', weights=ResNet18_Weights.SENTINEL2_RGB_MOCO)\n",
    "model = model.eval().to(accelerator)\n",
    "\n",
    "augs = K.AugmentationSequential(\n",
    "    K.Normalize(mean=datamodule.mean, std=datamodule.std, p=1.0)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def embed_standard(model, dataloader, accelerator, transforms=None):\n",
    "    x_all = []\n",
    "    y_all = []\n",
    "    for batch in tqdm(dataloader):\n",
    "        x = batch['image'].to(accelerator)\n",
    "        y = batch['label']\n",
    "        x = x.to(accelerator)\n",
    "        if transforms is not None:\n",
    "            x = transforms(x)\n",
    "\n",
    "        with torch.inference_mode():\n",
    "            embeddings = model.forward_features(x)\n",
    "            # global average pooling over the spatial dims\n",
    "            embeddings = torch.mean(embeddings, dim=(-2, -1))\n",
    "        x_all.append(embeddings.cpu().numpy())\n",
    "        y_all.append(y.numpy())\n",
    "    x_all = np.concatenate(x_all, axis=0)\n",
    "    y_all = np.concatenate(y_all, axis=0)\n",
    "    return x_all, y_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train, y_train = embed_standard(model, train_dl, accelerator, transforms=augs)\n",
    "x_val, y_val = embed_standard(model, val_dl, accelerator, transforms=augs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit/evaluate k-NN on ResNet-18 embeddings for a side-by-side comparison.\n",
    "\n",
    "knn_model = KNeighborsClassifier(n_neighbors=5)\n",
    "knn_model.fit(x_train, y_train)\n",
    "y_pred = knn_model.predict(x_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    classification_report(\n",
    "        y_val, y_pred, digits=2, target_names=class_names, zero_division=0\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = PCA(n_components=2, whiten=True)\n",
    "x_reduced = pca.fit_transform(x_train)\n",
    "print(f'Explained variance ratio: {pca.explained_variance_ratio_.sum()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(7, 7))\n",
    "ax = plt.gca()\n",
    "scatter = plt.scatter(\n",
    "    x_reduced[:, 0], x_reduced[:, 1], c=y_train, cmap='tab10', s=4, alpha=0.5\n",
    ")\n",
    "handles, _ = scatter.legend_elements()\n",
    "labels = [class_names[i] for i in range(10)]\n",
    "plt.legend(handles, labels, title='Classes', loc='best')\n",
    "plt.title('PCA of EuroSAT training set embeddings')\n",
    "plt.xlim([-1.5, 2])\n",
    "plt.ylim([-1, 2])\n",
    "plt.axis('off')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "plt.close()"
   ]
  }
 ],
 "metadata": {
  "execution": {
   "timeout": 1200
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}