{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6557fcc0cff9bb05",
   "metadata": {},
   "source": [
    "# Custom Feature Extractor with RootSIFT\n",
    "\n",
    "> [!NOTE]\n",
    "> Training the k-means model on every training image takes quite a bit of time. `IMAGE_STEP` keeps every `IMAGE_STEP`-th training image to train the embedder. Set it to `1` if you want to use all images.\n",
    "\n",
    "This notebook demonstrates how to write your own feature extractor by inheriting from `FeatureExtractorBase`. For this, a `RootSIFT` implementation on top of the `SIFT` implementation of `OpenCV` is used s a feature extractor instead of the built-in `RootSIFT` module, offering significantly more speed, at the cost of an `opencv` installation. It is then used to train a `VLADEmbedder` on the `Oxford Flowers` dataset and to compare two images."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "591180042dae7da9",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15d347c48dedfd91",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import torch\n",
    "\n",
    "from pyvisim.base import FeatureExtractorBase\n",
    "from pyvisim.classic import 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": "aa6a4cdcb085f5c5",
   "metadata": {},
   "source": [
    "## Hyperparameters\n",
    "\n",
    "`NUM_CLUSTERS` is the number of visual words of the VLAD vocabulary, `DIM_REDUCTION_FACTOR` reduces the dimension of the descriptors by half using `PCA` before the vocabulary is learned, and `NUM_FEATURES` is the maximum number of keypoints `SIFT` detects per image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cef394a02510c7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLUSTERS = 32\n",
    "DIM_REDUCTION_FACTOR = 2\n",
    "IMAGE_STEP = 6\n",
    "NUM_FEATURES = 500"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0f1680bee500135",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aed8ed960074655",
   "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()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99326864ad0e06de",
   "metadata": {},
   "source": [
    "## Declare the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03ec11747a00b4f9",
   "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": "4bb2201540bf7c37",
   "metadata": {},
   "source": [
    "### Plot some images from the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde982154fb962a4",
   "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": "4532d958fe30fe13",
   "metadata": {},
   "source": [
    "## Write the RootSIFT feature extractor\n",
    "\n",
    "A feature extractor maps one image to an `(N, D)` array of local descriptors, which an embedder then aggregates into a single vector. Following methods must be implemented:\n",
    "\n",
    "- `output_dim`: the dimension `D` of one descriptor.\n",
    "- `__call__`: takes one image and returns its descriptors as a `float32` array of shape `(N, output_dim)`.\n",
    "\n",
    "`__call__` accepts any `MatLike` image (a NumPy array, a torch tensor or an array-like object) in the layout given by `dims` and the value range given by `value_range`. `to_single_image` converts it into the canonical `uint8` array of shape `(H, W, C)`, or `(H, W)` for a grayscale input.\n",
    "\n",
    "`RootSIFT` [2] starts from the `SIFT` [1] descriptors computed by `OpenCV`, 128 values per keypoint. Each descriptor is divided by its L1 norm, and then every value is replaced by its square root. Comparing two `RootSIFT` descriptors with the Euclidean distance is then the same as comparing the original `SIFT` descriptors with the Hellinger kernel, which suits histograms such as `SIFT` better [2]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aa348e328ee76f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class OpenCV_RootSIFT(FeatureExtractorBase):\n",
    "    \"\"\"\n",
    "    Scale-Invariant Feature Transform with Hellinger kernel (RootSIFT) feature extractor.\n",
    "\n",
    "    :param n_features: Maximum number of keypoints to detect per image.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, n_features: int = 500) -> None:\n",
    "        super().__init__()\n",
    "        self._sift = cv2.SIFT_create(nfeatures=n_features)\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": "markdown",
   "id": "9b8700f2cf6f472a",
   "metadata": {},
   "source": [
    "Let's check the extractor on one image. Every row of the output is the descriptor of one keypoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f395e72d4fecc1b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "extractor = OpenCV_RootSIFT(n_features=NUM_FEATURES)\n",
    "descriptors = extractor(train_dataset[0][0])\n",
    "print(\"Output dimension:\", extractor.output_dim)\n",
    "print(\"Descriptors shape:\", descriptors.shape)\n",
    "print(\"Descriptors dtype:\", descriptors.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08bd986b62491dde",
   "metadata": {},
   "source": [
    "## Declare the VLAD embedder\n",
    "\n",
    "The extractor is passed to the `VLADEmbedder` [3]. The embedder calls the extractor on every image, reduces the descriptors with `PCA` and clusters them into `NUM_CLUSTERS` visual words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8607142a0a45e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "vlad_embedder = VLADEmbedder(feature_extractor=extractor, n_clusters=NUM_CLUSTERS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "441a0fe2d6f4ef6e",
   "metadata": {},
   "source": [
    "The following cell trains the model from scratch on the training images. The dimension of the descriptors is reduced by half using `PCA` before the k-means model is trained. It might take quite a bit of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4da2a7a989f4780",
   "metadata": {},
   "outputs": [],
   "source": [
    "vlad_embedder.learn(\n",
    "    (train_dataset[i][0] for i in train_indices), dim_reduction_factor=DIM_REDUCTION_FACTOR\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8371009184c2f161",
   "metadata": {},
   "source": [
    "## Compare two images\n",
    "\n",
    "Now, we pick a reference image, a similar image and a dissimilar image from the validation set, on which the model is not trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "655a21f175c5a191",
   "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": "6a27a3f8e071e577",
   "metadata": {},
   "source": [
    "Now, we 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "813ea93d01061496",
   "metadata": {},
   "outputs": [],
   "source": [
    "score_similar = vlad_embedder.similarity_score(image_ref, image_similar).item()\n",
    "print(f\"Similarity score, similar pair: {score_similar:.4f}\")\n",
    "\n",
    "score_dissimilar = vlad_embedder.similarity_score(image_ref, image_dissimilar).item()\n",
    "print(f\"Similarity score, dissimilar pair: {score_dissimilar:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4270c9edb74778d",
   "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] Arandjelović, R., & Zisserman, A. (2013). All About VLAD. In 2013 IEEE\n",
    "Conference on Computer Vision and Pattern Recognition (CVPR), 1578-1585.\n",
    "https://doi.org/10.1109/CVPR.2013.207"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}