{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "25d52ed7503c2a1f",
   "metadata": {},
   "source": [
    "# Classical Embedding algorithms: `VLAD` and `Fisher Vectors`\n",
    "\n",
    "> [!NOTE]\n",
    "> Extracting the descriptors and training the clustering models on every training image takes quite a bit of time. `IMAGE_STEP` keeps every `IMAGE_STEP`-th training image to train the embedders. Set it to `1` if you want to use all images.\n",
    "\n",
    "Vector of Locally Aggregated Descriptors (VLAD) [3] and Fisher Vectors [4] turn the local descriptors of an image into a single vector of fixed size. Both first learn a visual vocabulary from the descriptors of the training images: VLAD clusters them with **k-means**, and the Fisher Vector fits a **Gaussian mixture model** to them.\n",
    "\n",
    "This notebook trains both embedders on `RootSIFT` descriptors of the `Oxford Flowers` dataset in two ways:\n",
    "\n",
    "1. With `learn`, which extracts the descriptors, fits the `PCA` and fits the clustering model in one call. This uses this library's built-in clustering and PCA implementations and hence spare a `scikit-learn` installation.\n",
    "2. With `scikit-learn` estimators: if you need extensive training capabilities of `scikit-learn`, you can train your own `PCA` and clustering models first, and then transfer their parameters to the embedders with `load_clustering_model_from_sklearn` and `load_pca_model_from_sklearn`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a34643ae72491537",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c5f0c991d426e16",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from matplotlib import pyplot as plt\n",
    "from sklearn.cluster import KMeans\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.mixture import GaussianMixture\n",
    "\n",
    "from pyvisim.base import FeatureExtractorBase\n",
    "from pyvisim.classic import FisherVectorEmbedder, VLADEmbedder\n",
    "from pyvisim.datasets import OxfordFlowerDataset\n",
    "from pyvisim.typing import Float32NumpyArray, MatLike\n",
    "from pyvisim.utils.image_utils import to_single_image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a66a68bbe486928",
   "metadata": {},
   "source": [
    "## Hyperparameters\n",
    "\n",
    "`NUM_CLUSTERS` is the size of the visual vocabulary: the number of k-means clusters for VLAD and the number of Gaussian components for the Fisher Vector. `DIM_REDUCTION_FACTOR` halves the dimension of the descriptors with `PCA` before the vocabulary is learned, and `SEED` makes the training reproducible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "650cc6d4d27c8969",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLUSTERS = 32\n",
    "DIM_REDUCTION_FACTOR = 2\n",
    "IMAGE_STEP = 6\n",
    "SEED = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cd019d0492cf50e",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f8e5eeab4aa2782",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_image(image: np.ndarray | torch.Tensor, title: str = \"Image\") -> None:\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    if isinstance(image, torch.Tensor):\n",
    "        image = image.detach().cpu()\n",
    "        if image.ndim == 3:\n",
    "            image = image.permute(1, 2, 0)\n",
    "        image = image.numpy()\n",
    "    plt.imshow(image)\n",
    "    plt.axis(\"off\")\n",
    "    plt.title(title)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def print_similarity_scores(\n",
    "    embedder: VLADEmbedder | FisherVectorEmbedder,\n",
    "    reference: MatLike,\n",
    "    similar: MatLike,\n",
    "    dissimilar: MatLike,\n",
    ") -> None:\n",
    "    score_similar = embedder.similarity_score(reference, similar).item()\n",
    "    score_dissimilar = embedder.similarity_score(reference, dissimilar).item()\n",
    "    print(f\"Similarity score, similar pair: {score_similar:.4f}\")\n",
    "    print(f\"Similarity score, dissimilar pair: {score_dissimilar:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ebce99e32949753",
   "metadata": {},
   "source": [
    "## Declare the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "069bcb63bda5ab5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = OxfordFlowerDataset(purpose=\"train\")\n",
    "val_dataset = OxfordFlowerDataset(purpose=\"validation\")\n",
    "print(\"Number of images in the dataset:\", len(train_dataset))\n",
    "\n",
    "train_indices = range(0, len(train_dataset), IMAGE_STEP)\n",
    "print(\"Number of images used for training:\", len(train_indices))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dccb407fad7186c",
   "metadata": {},
   "source": [
    "### Plot some images from the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd702b353409aec",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(3):\n",
    "    img, label, _ = train_dataset[i]\n",
    "    print(\"Image size:\", img.shape)\n",
    "    plot_image(img, title=f\"Label: {label}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ea5052d1155cdf7",
   "metadata": {},
   "source": [
    "## Declare the RootSIFT feature extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d13c0d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyvisim.features import RootSIFT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38272c1f",
   "metadata": {},
   "source": [
    "> [!NOTE]\n",
    "> The `RootSIFT` implementation of this library is written in Python, and hence not as performant as `OpenCV's`. If you need more speed, create your own feature extractor (see class `OpenCV_RootSIFT` below) using `OpenCV`. This is further covered in the [RootSIFT tutorial](custom_feature_extractor_with_rootsift.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56f3105d8958d8dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyvisim.lazy_import import OptionalImport\n",
    "\n",
    "with OptionalImport(err_msg = \"OpenCV is not installed. Please install it with `pip install opencv-python-headless`\") as cv2_import:\n",
    "    import cv2\n",
    "\n",
    "\n",
    "class OpenCV_RootSIFT(FeatureExtractorBase):\n",
    "    \"\"\"\n",
    "    Scale-Invariant Feature Transform with Hellinger kernel (RootSIFT) feature extractor.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__()\n",
    "        cv2_import.check()\n",
    "        self._sift = cv2.SIFT_create()\n",
    "\n",
    "    @property\n",
    "    def output_dim(self) -> int:\n",
    "        return self._sift.descriptorSize()\n",
    "\n",
    "    def __call__(\n",
    "        self,\n",
    "        image: MatLike,\n",
    "        /,\n",
    "        *,\n",
    "        dims: str = \"HWC\",\n",
    "        value_range: tuple[float, float] = (0.0, 255.0),\n",
    "    ) -> Float32NumpyArray:\n",
    "        canonical = to_single_image(image, dims=dims, value_range=value_range)\n",
    "        grayscale = (\n",
    "            cv2.cvtColor(canonical, cv2.COLOR_RGB2GRAY)\n",
    "            if canonical.ndim == 3\n",
    "            else canonical\n",
    "        )\n",
    "        _, descriptors = self._sift.detectAndCompute(grayscale, None)\n",
    "        if descriptors is None:\n",
    "            return np.zeros((0, self.output_dim), dtype=np.float32)\n",
    "        descriptors = descriptors / (descriptors.sum(axis=1, keepdims=True) + 1e-7)\n",
    "        return np.sqrt(descriptors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78723bac7771484a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# To speed up CI on GitHub Actions, OpenCV_RootSIFT is used instead of the built-in RootSIFT implementation. Comment in \n",
    "# the line below to use the built-in RootSIFT, which saves the dependency on OpenCV, but is slower.\n",
    "# extractor = RootSIFT()\n",
    "\n",
    "extractor = OpenCV_RootSIFT()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc21a2d59f78a427",
   "metadata": {},
   "source": [
    "## Pick the images to compare\n",
    "\n",
    "Every trained embedder scores the same three images from the validation set, on which none of them is trained: a reference image, a similar image of the same class and a dissimilar image of another class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e8161380ec2e498",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_ref, label_ref, _ = val_dataset[2]\n",
    "image_similar, label_similar, _ = val_dataset[3]\n",
    "image_dissimilar, label_dissimilar, _ = val_dataset[100]\n",
    "plot_image(image_ref, title=f\"Reference Image. Label: {label_ref}\")\n",
    "plot_image(image_similar, title=f\"Similar Image. Label: {label_similar}\")\n",
    "plot_image(image_dissimilar, title=f\"Dissimilar Image. Label: {label_dissimilar}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8458eb2951061608",
   "metadata": {},
   "source": [
    "## VLAD\n",
    "\n",
    "VLAD assigns every descriptor of an image to its nearest k-means centroid and sums the residuals, the differences between the descriptors and their centroid, per cluster. The sums of all clusters are concatenated into one vector of size `K * D`, where `K` equals `NUM_CLUSTERS` and `D` is the dimension of the descriptors after `PCA`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaf55f524530f64e",
   "metadata": {},
   "source": [
    "### Train with `learn`\n",
    "\n",
    "`learn` extracts the descriptors of the training images, fits a `PCA` that keeps `output_dim // DIM_REDUCTION_FACTOR` components and fits the k-means model on the projected descriptors. The options of the k-means model go into `kmeans_params`. The following cell might take quite a bit of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97ba2914bd2a74f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "vlad_embedder = VLADEmbedder(\n",
    "    feature_extractor=extractor, n_clusters=NUM_CLUSTERS, kmeans_params={\"rng\": SEED}\n",
    ")\n",
    "vlad_embedder.learn(\n",
    "    (train_dataset[i][0] for i in train_indices), dim_reduction_factor=DIM_REDUCTION_FACTOR\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a680b960192388",
   "metadata": {},
   "source": [
    "### Train with scikit-learn\n",
    "\n",
    "To use `scikit-learn`, you extract the descriptors yourself. The extractor returns one `(N, D)` array per image, and `np.vstack` stacks them into a single matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed6a640e2f32d7d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "descriptors = np.vstack([extractor(train_dataset[i][0]) for i in train_indices])\n",
    "print(\"Descriptors shape:\", descriptors.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34467a4afb8dcfe4",
   "metadata": {},
   "source": [
    "Next comes the `PCA`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f93888495b93d523",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = PCA(\n",
    "    n_components=extractor.output_dim // DIM_REDUCTION_FACTOR,\n",
    "    whiten=True,\n",
    "    svd_solver=\"randomized\",\n",
    "    random_state=SEED,\n",
    ")\n",
    "projected_descriptors = pca.fit_transform(descriptors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2ee41f86a2d831b",
   "metadata": {},
   "source": [
    "The k-means model is fitted on the projected descriptors. `max_iter` stops the training after 50 iterations. The built-in k-means has no such limit and iterates until the change in distortion falls below `thresh`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "942865b72b04996e",
   "metadata": {},
   "outputs": [],
   "source": [
    "kmeans = KMeans(\n",
    "    n_clusters=NUM_CLUSTERS, max_iter=50, algorithm=\"lloyd\", random_state=SEED\n",
    ")\n",
    "kmeans.fit(projected_descriptors)\n",
    "print(\"Number of iterations:\", kmeans.n_iter_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7286d307b9a569d",
   "metadata": {},
   "source": [
    "Now, the fitted estimators go into a new `VLADEmbedder`. `load_pca_model_from_sklearn` and `load_clustering_model_from_sklearn` convert them into the models of this library together with everything they learned. The number of clusters comes from the k-means model, so `n_clusters` is left out.\n",
    "\n",
    "The `PCA` has to be set first. The embedder checks that the input dimension of the k-means model matches the output dimension of the `PCA`, and without a `PCA` it compares it with the dimension of the raw descriptors, which fails.\n",
    "\n",
    "> [!NOTE]\n",
    "> Fit the estimators before you pass them in. An unfitted estimator only passes on the options that the built-in models also have, and `learn` then trains the built-in models. `scikit-learn` options such as `max_iter`, `algorithm` or `svd_solver=\"randomized\"` are lost that way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "861d89ba88ed9aea",
   "metadata": {},
   "outputs": [],
   "source": [
    "vlad_embedder_sklearn = VLADEmbedder(feature_extractor=extractor)\n",
    "vlad_embedder_sklearn.load_pca_model_from_sklearn(pca)\n",
    "vlad_embedder_sklearn.load_clustering_model_from_sklearn(kmeans)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff4095e8e2031c5e",
   "metadata": {},
   "source": [
    "After the conversion, the embedder no longer needs the `scikit-learn` estimators. It is saved with `save_to_disk` and loaded with `load_from_disk` like an embedder trained with `learn`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f568475e56d970eb",
   "metadata": {},
   "source": [
    "### Compare the images\n",
    "\n",
    "Both VLAD embedders compare the reference image with the other two. `cosine similarity` is used in this case, so the score lies in `[-1, 1]` and a higher score means more similar. The two embedders were trained with different options, so their scores differ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7649aee1912d372b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"VLAD trained with learn:\")\n",
    "print_similarity_scores(vlad_embedder, image_ref, image_similar, image_dissimilar)\n",
    "\n",
    "print(\"\\nVLAD trained with scikit-learn:\")\n",
    "print_similarity_scores(\n",
    "    vlad_embedder_sklearn, image_ref, image_similar, image_dissimilar\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f939cd99d4021a2",
   "metadata": {},
   "source": [
    "## Fisher Vector\n",
    "\n",
    "The Fisher Vector fits a Gaussian mixture model with a diagonal covariance matrix per component instead of k-means. Every descriptor is assigned to all components, weighted by its posterior probabilities, and the vector holds the gradients of the log-likelihood with respect to the mixture weights, the means and the variances. With `K` components and descriptors of dimension `D`, it has `2 * K * D + K` values, about twice as many as a VLAD vector."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43eef2b28567ba33",
   "metadata": {},
   "source": [
    "### Train with `learn`\n",
    "\n",
    "Training works the same way as for VLAD. `n_components` sets the number of Gaussian components, and further options of the Gaussian mixture model go into `gmm_params`. The following cell might take quite a bit of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07497a6b7cbbefa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_embedder = FisherVectorEmbedder(\n",
    "    feature_extractor=extractor, n_components=NUM_CLUSTERS, gmm_params={\"rng\": SEED}\n",
    ")\n",
    "fisher_embedder.learn(\n",
    "    (train_dataset[i][0] for i in train_indices), dim_reduction_factor=DIM_REDUCTION_FACTOR\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c30727bd8d5bc1fb",
   "metadata": {},
   "source": [
    "### Train with scikit-learn\n",
    "\n",
    "The descriptors and the fitted `PCA` from the VLAD section are reused. Keep in mind that `covariance_type=\"diag\"` is required, since the Fisher Vector of this library only supports diagonal covariances.\n",
    "\n",
    "`init_params=\"k-means++\"` starts the EM algorithm from k-means++ seeds, which is faster than the full k-means run of the default `\"kmeans\"`. Other values of `init_params` cannot be converted. `verbose=2` prints the progress of the EM algorithm every `verbose_interval` iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc22d9470f7732d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmm = GaussianMixture(\n",
    "    n_components=NUM_CLUSTERS,\n",
    "    covariance_type=\"diag\",\n",
    "    init_params=\"k-means++\",\n",
    "    verbose=2,\n",
    "    verbose_interval=10,\n",
    "    random_state=SEED,\n",
    ")\n",
    "gmm.fit(projected_descriptors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "115b7538a552e2ef",
   "metadata": {},
   "source": [
    "As for VLAD, the `PCA` goes in before the Gaussian mixture model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a986bd58b1843c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_embedder_sklearn = FisherVectorEmbedder(feature_extractor=extractor)\n",
    "fisher_embedder_sklearn.load_pca_model_from_sklearn(pca)\n",
    "fisher_embedder_sklearn.load_clustering_model_from_sklearn(gmm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28f4ad02c2a81d92",
   "metadata": {},
   "source": [
    "### Compare the images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8f309b7732f049d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Fisher Vector trained with learn:\")\n",
    "print_similarity_scores(fisher_embedder, image_ref, image_similar, image_dissimilar)\n",
    "\n",
    "print(\"\\nFisher Vector trained with scikit-learn:\")\n",
    "print_similarity_scores(\n",
    "    fisher_embedder_sklearn, image_ref, image_similar, image_dissimilar\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2162541d7f855c1",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "[1] Lowe, D. G. (2004). Distinctive image features from scale-invariant\n",
    "keypoints. International Journal of Computer Vision, 60(2), 91-110.\n",
    "https://doi.org/10.1023/B:VISI.0000029664.99615.94\n",
    "\n",
    "[2] Arandjelović, R., & Zisserman, A. (2012). Three things everyone should know\n",
    "to improve object retrieval. In 2012 IEEE Conference on Computer Vision and\n",
    "Pattern Recognition (CVPR), 2911-2918. https://doi.org/10.1109/CVPR.2012.6248018\n",
    "\n",
    "[3] Jégou, H., Douze, M., Schmid, C., & Pérez, P. (2010). Aggregating local\n",
    "descriptors into a compact image representation. In 2010 IEEE Computer Society\n",
    "Conference on Computer Vision and Pattern Recognition (CVPR), 3304-3311.\n",
    "https://doi.org/10.1109/CVPR.2010.5540039\n",
    "\n",
    "[4] Jégou, H., Perronnin, F., Douze, M., Sánchez, J., Pérez, P., & Schmid, C.\n",
    "(2012). Aggregating local image descriptors into compact codes. IEEE\n",
    "Transactions on Pattern Analysis and Machine Intelligence, 34(9), 1704-1716.\n",
    "https://doi.org/10.1109/TPAMI.2011.235"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}