{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5b0885d297b0ea1d",
   "metadata": {},
   "source": [
    "# Custom Feature Extractor with ORB\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, `ORB` from `OpenCV` is implemented as a feature extractor, which is then used to train a `VLADEmbedder` on the `Oxford Flowers` dataset and to compare two images."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c902a80a20a8187",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10cb0225e9be91c9",
   "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.features._utils import _to_single_image\n",
    "from pyvisim.typing import Float32NumpyArray, MatLike"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac16c6f9541a3572",
   "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 `ORB` detects per image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d826ff955eb60e34",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLUSTERS = 32\n",
    "DIM_REDUCTION_FACTOR = 2\n",
    "IMAGE_STEP = 4\n",
    "NUM_FEATURES = 500"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47832892",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1a2bd22",
   "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": "c7cc199ad70a8b54",
   "metadata": {},
   "source": [
    "## Declare the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0615caa9b51c5cab",
   "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": "bad01e0f842d5aba",
   "metadata": {},
   "source": [
    "### Plot some images from the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deaa58ef77c93ab9",
   "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": "354699a8d35cd39e",
   "metadata": {},
   "source": [
    "## Write the ORB 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. `FeatureExtractorBase` asks for exactly two things from a subclass:\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, so the extractor only has to deal with one input format. This is also what the built-in extractors such as `SIFT` do.\n",
    "\n",
    "`ORB` [1] is a binary descriptor: each keypoint is described by 32 bytes, that is 256 bits. `k-means` needs real-valued vectors, so the bits are unpacked into a vector of 256 zeros and ones. If `ORB` finds no keypoints in an image, `OpenCV` returns `None`, in which case an empty `(0, output_dim)` array is returned instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a731050821b8972",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ORB(FeatureExtractorBase):\n",
    "    \"\"\"\n",
    "    Oriented FAST and Rotated BRIEF (ORB) 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._orb = cv2.ORB_create(nfeatures=n_features)\n",
    "\n",
    "    @property\n",
    "    def output_dim(self) -> int:\n",
    "        # OpenCV stores each descriptor as bytes, one bit per BRIEF test.\n",
    "        return 8 * self._orb.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._orb.detectAndCompute(grayscale, None)\n",
    "        if descriptors is None:\n",
    "            return np.zeros((0, self.output_dim), dtype=np.float32)\n",
    "        return np.unpackbits(descriptors, axis=1).astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "248bdbce8edc63d5",
   "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": "5b21140f5e517a18",
   "metadata": {},
   "outputs": [],
   "source": [
    "extractor = ORB(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": "fd72acc628888d14",
   "metadata": {},
   "source": [
    "## Declare the VLAD embedder\n",
    "\n",
    "The extractor is passed to the `VLADEmbedder` like any built-in one. 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": "f6e785c8582c2a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vlad_embedder = VLADEmbedder(feature_extractor=extractor, n_clusters=NUM_CLUSTERS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db0a0403231c2956",
   "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": "566cad3e49edf035",
   "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": "dbff423d9b5ea145",
   "metadata": {},
   "source": [
    "## Compare two images\n",
    "\n",
    "Now, we will pick one image from the training set and one from the validation set, on which the model is not yet trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b989c5e0b3fee37f",
   "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": "08b1dfa547e2e363",
   "metadata": {},
   "source": [
    "Now, we compare the two images. `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": "a2d62683db004d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "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": "6c188d10dd08233d",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "[1] Rublee, E., Rabaud, V., Konolige, K., & Bradski, G. (2011). ORB: An\n",
    "efficient alternative to SIFT or SURF. In 2011 International Conference on\n",
    "Computer Vision (ICCV), 2564-2571. https://doi.org/10.1109/ICCV.2011.6126544\n",
    "\n",
    "[2] 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
}