{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "69119ca9c7123f0c",
   "metadata": {},
   "source": [
    "# **Pipeline with Deep Features**\n",
    "\n",
    "In this notebook, we'll:\n",
    "1. Load (or extract) deep features from a set of images (Oxford Flowers or any dataset).\n",
    "2. Train a single KMeans and GMM model (both with 32 clusters/components).\n",
    "3. Set up a VLAD embedder using the KMeans model.\n",
    "4. Set up a Fisher Vector embedder using the GMM model.\n",
    "5. Use each embedder separately for retrieval, then combine them in a pipeline, and compare retrieval metrics.\n",
    "\n",
    "For more detail, see the paper below.\n",
    "\n",
    "## Reference:\n",
    "[1]Weixia Zhang, Jia Yan, Wenxuan Shi, Tianpeng Feng, and Dexiang Deng, \"Refining Deep Convolutional Features for Improving Fine-Grained Image Recognition,\" EURASIP Journal on Image and Video Processing, 2017."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cfeb663c836d2a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Any\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import torch\n",
    "\n",
    "from pyvisim.distance import cosine_similarity\n",
    "from pyvisim.classic import FisherVectorEmbedder, Pipeline, VLADEmbedder\n",
    "from pyvisim.datasets import OxfordFlowerDataset  # or any dataset you have\n",
    "from pyvisim.neural_networks.features import DeepConvFeature\n",
    "from pyvisim.retrieval.image_store import InMemoryImageEmbeddingStore\n",
    "from pyvisim.typing import NumpyArray, UInt8NumpyArray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e9ea01c",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7044caca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_image(image: UInt8NumpyArray | torch.Tensor, title: str = \"Image\") -> None:\n",
    "    \"\"\"\n",
    "    Plot a single image.\n",
    "\n",
    "    :param image: Image as a NumPy array (H, W, C) or torch tensor (C, H, W)\n",
    "    :param title: Title of the plot\n",
    "    \"\"\"\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 plot_and_save_heatmap(\n",
    "    matrix: list[Any] | NumpyArray | torch.Tensor,\n",
    "    figsize: tuple[int, int] | None = None,\n",
    "    x_tick_labels: list[str] | None = None,\n",
    "    y_tick_labels: list[str] | None = None,\n",
    "    cbar_kws: dict[str, str] | None = None,\n",
    "    title: str = \"Heatmap\",\n",
    "    x_label: str = \"X Axis\",\n",
    "    y_label: str = \"Y Axis\",\n",
    "    show: bool = True,\n",
    "    save_fig_path: str | None = None,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Plot a heatmap using the specified matrix.\n",
    "\n",
    "    :param matrix: matrix\n",
    "    :param figsize: figure size\n",
    "    :param x_tick_labels: x-axis tick labels\n",
    "    :param y_tick_labels: y-axis tick labels\n",
    "    :param cbar_kws: colorbar keyword arguments\n",
    "    :param title: title of the plot\n",
    "    :param x_label: x-axis label\n",
    "    :param y_label: y-axis label\n",
    "    :param show: whether to display the plot\n",
    "    :param save_fig_path: Path to save the figure\n",
    "    \"\"\"\n",
    "    if isinstance(matrix, list):\n",
    "        matrix = np.array(matrix)\n",
    "    elif isinstance(matrix, torch.Tensor):\n",
    "        matrix = matrix.detach().cpu().numpy()\n",
    "\n",
    "    figsize = (\n",
    "        (\n",
    "            int(matrix.shape[1] * 0.7),\n",
    "            int(matrix.shape[0] * 0.7),\n",
    "        )\n",
    "        if figsize is None\n",
    "        else figsize\n",
    "    )\n",
    "    plt.figure(figsize=figsize)\n",
    "    sns.heatmap(\n",
    "        matrix,\n",
    "        annot=True,\n",
    "        fmt=\".2f\",\n",
    "        cmap=\"viridis\",\n",
    "        xticklabels=x_tick_labels if x_tick_labels else list(range(matrix.shape[1])),\n",
    "        yticklabels=y_tick_labels if y_tick_labels else list(range(matrix.shape[0])),\n",
    "        cbar_kws=cbar_kws if cbar_kws else {\"label\": \"value\"},\n",
    "    )\n",
    "    plt.title(title)\n",
    "    plt.xlabel(x_label)\n",
    "    plt.ylabel(y_label)\n",
    "    if save_fig_path:\n",
    "        plt.savefig(save_fig_path)\n",
    "    if show:\n",
    "        plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad9ec835f509c70",
   "metadata": {},
   "source": [
    "## **1. Load Dataset and Extract Deep Features**\n",
    "We'll load a subset of images and collect their deep descriptors.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f872152ef51fe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = OxfordFlowerDataset(\n",
    "    purpose=\"train\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75b146e7eeb847f4",
   "metadata": {},
   "source": [
    "### Let's define our deep extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "557227f17fb60fe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "extractor = DeepConvFeature(\n",
    "    backbone=\"resnet18\",\n",
    "    layer_index=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e489b81738be233d",
   "metadata": {},
   "source": [
    "## **2. Initialize Embedders: VLAD and Fisher Vector**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e49930fd24d6d0dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "vlad_embedder = VLADEmbedder(\n",
    "    feature_extractor=extractor,\n",
    "    n_clusters=32,\n",
    "    power_norm_weight=1,\n",
    ")\n",
    "\n",
    "fisher_embedder = FisherVectorEmbedder(\n",
    "    feature_extractor=extractor,\n",
    "    n_components=32,\n",
    "    power_norm_weight=0.5,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4289642",
   "metadata": {},
   "source": [
    "### Train the KMeans and GMM models\n",
    "\n",
    "Both models are trained on the deep features of every `IMAGE_STEP`-th training image. The training split is sorted by class, so the step keeps all 102 classes represented. Set it to `1` to train on the whole training split, which takes quite a bit of time. For the Fisher Vector, the feature dimension is reduced by half using `PCA` first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7d265d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_STEP = 12\n",
    "train_indices = range(0, len(train_dataset), IMAGE_STEP)\n",
    "print(f\"Training on {len(train_indices)} images\")\n",
    "\n",
    "vlad_embedder.learn(train_dataset[i][0] for i in train_indices)\n",
    "fisher_embedder.learn(\n",
    "    (train_dataset[i][0] for i in train_indices),\n",
    "    dim_reduction_factor=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "639a15e3c3c45e76",
   "metadata": {},
   "source": [
    "## **4. Compute vectors**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b61cac7dd016af7f",
   "metadata": {},
   "source": [
    "### Select images for comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ce04dca76bc5817",
   "metadata": {},
   "outputs": [],
   "source": [
    "image1, *_ = train_dataset[0]\n",
    "image2, *_ = train_dataset[1]\n",
    "plot_image(image1)\n",
    "plot_image(image2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9f2f8d24dc940f5",
   "metadata": {},
   "source": [
    "### Compute VLAD and Fisher Vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b02af153a024ba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_vector_1 = fisher_embedder.embed(image1)\n",
    "vlad_vector_1 = vlad_embedder.embed(image1)\n",
    "print(\n",
    "    f\"Shape of Fisher Vector: {fisher_vector_1.shape}, Shape of VLAD Vector: {vlad_vector_1.shape}\"\n",
    ")\n",
    "\n",
    "fisher_vector_2 = fisher_embedder.embed(image2)\n",
    "vlad_vector_2 = vlad_embedder.embed(image2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8025b2b65d617b6c",
   "metadata": {},
   "source": [
    "### Compute similarity scores using VLAD and Fisher Vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "699cedcdaf7b4736",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Cosine similarity using VLAD:\", vlad_embedder.similarity_score(image1, image2))\n",
    "print(\n",
    "    \"Cosine similarity using Fisher Vectors:\",\n",
    "    fisher_embedder.similarity_score(image1, image2),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8371fc78e6d457",
   "metadata": {},
   "source": [
    "## Use a Pipeline\n",
    "\n",
    "We can combine both steps above by using the pipeline class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "336b8eef79a64d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = Pipeline([vlad_embedder, fisher_embedder])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d0d10ccfd6c7e5a",
   "metadata": {},
   "source": [
    "### Compare images using the pipeline\n",
    "\n",
    "The result is the similarity score of the concatenated VLAD and Fisher vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d1d52058a2cb366",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_score = pipeline.similarity_score(image1, image2)\n",
    "print(\"Scores using the pipeline:\", sim_score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6b6fd4363d18589",
   "metadata": {},
   "source": [
    "This is equal to:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd795ce6c10526e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_vector_1 = np.hstack((vlad_vector_1, fisher_vector_1))\n",
    "combined_vector_2 = np.hstack((vlad_vector_2, fisher_vector_2))\n",
    "print(\n",
    "    \"Cosine similarity using concatenated vectors:\",\n",
    "    cosine_similarity(combined_vector_1, combined_vector_2),\n",
    ")  # Cosine similarity was chosen as `similarity_func` for the embedders"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0e6d96bd0c80dc5",
   "metadata": {},
   "source": [
    "You can also compare two batches of images at once. The shape is the similarity matrix of size (batch_1_size, batch_2_size)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74659117a5ca28f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_1 = (train_dataset[i][0] for i in range(5))\n",
    "batch_2 = (train_dataset[i][0] for i in range(5, 15))\n",
    "similarity_matrix = pipeline.similarity_score(batch_1, batch_2)\n",
    "plot_and_save_heatmap(\n",
    "    similarity_matrix,\n",
    "    x_tick_labels=[f\"Image {i}\" for i in range(10)],\n",
    "    y_tick_labels=[f\"Image {i}\" for i in range(5)],\n",
    "    figsize=(10, 5),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "534a7c2e1acb0523",
   "metadata": {},
   "source": [
    "You can call the `embed` method, just like with `VLADEmbedder` and `FisherVectorEmbedder`. Each image is embedded using the embedders in the pipeline, and the end result is the concatenation of the vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2188eb6c839c9e41",
   "metadata": {},
   "outputs": [],
   "source": [
    "images_1_5 = [train_dataset[i][0] for i in range(5)]\n",
    "embedding = pipeline.embed(images_1_5)\n",
    "print(\"Shape of embedded batch:\", embedding.shape)  # (num_images, dim_vlad + dim_fisher)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eafffddc914c9e9",
   "metadata": {},
   "source": [
    "### Build an image store\n",
    "\n",
    "The store embeds the images with the pipeline. Its paths are the image paths and its embeddings are the concatenated VLAD and Fisher vectors, L2-normalised since the store is built in the `cosine` space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e303f792f98f9a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_paths = train_dataset.image_paths[:5]\n",
    "store = InMemoryImageEmbeddingStore(image_paths=image_paths, embedder=pipeline)\n",
    "store.build_store()\n",
    "{\n",
    "    os.path.basename(path): vector\n",
    "    for path, vector in zip(store.paths, store.embeddings, strict=True)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b07f6ef2a0c690",
   "metadata": {},
   "source": [
    "Print the pipeline to see the embedders it contains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bae732d343e92a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(pipeline)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}