{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bf8480748f688a15",
   "metadata": {},
   "source": [
    "# **Computing `Mean Average Precision (mAP)` and `Top-k Accuracy` for our Retrieval System**\n",
    "We'll use every `IMAGE_STEP`-th image in the `validation + test dataset` as a query. For each query:\n",
    "1. Retrieve all images, rank them by similarity.\n",
    "2. Compute average precision for each query.\n",
    "3. Take the mean across all queries => mAP. This takes into consideration the ranking of the images.\n",
    "4. We will evaluate top 1 accuracy and top-k accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a6d104abe78a29b",
   "metadata": {},
   "source": [
    "## **1. Import Necessary Libraries**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ded369c2ee9dacc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from pyvisim.datasets import OxfordFlowerDataset\n",
    "from pyvisim.eval import top_k_accuracy, top_k_map\n",
    "from pyvisim.neural_networks import ClipEmbedder\n",
    "from pyvisim.retrieval.image_store import InMemoryImageEmbeddingStore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78f141f1",
   "metadata": {},
   "source": [
    "### Hyperparameters\n",
    "\n",
    "> [!NOTE]\n",
    ">\n",
    "> `IMAGE_STEP` keeps every `IMAGE_STEP`-th training image in the image store and every `IMAGE_STEP`-th validation and test image as a query. Set to `1` if you want to use all images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b381e16",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_STEP = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77cbbd65",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae7260bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_and_save_barplot(\n",
    "    data: dict[str, list[float]],\n",
    "    bar_labels: list[str],\n",
    "    title: str = \"Barplot\",\n",
    "    xlabel: str = \"X-axis\",\n",
    "    ylabel: str = \"Y-axis\",\n",
    "    save_path: str | None = None,\n",
    "    show: bool = True,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Plot and save a barplot.\n",
    "\n",
    "    :param data: Dictionary containing data to plot.\n",
    "    :param bar_labels: Labels that will be displayed in the legend.\n",
    "    :param title: Title of the plot.\n",
    "    :param xlabel: Label for the x-axis.\n",
    "    :param ylabel: Label for the y-axis.\n",
    "    :param save_path: Path to save the plot image. If None, plot is not saved.\n",
    "    :param show: Whether to display the plot.\n",
    "    \"\"\"\n",
    "    x_labels = list(data.keys())\n",
    "    values = list(data.values())\n",
    "    num_groups = len(values[0])\n",
    "\n",
    "    if not all(len(v) == num_groups for v in values):\n",
    "        raise ValueError(\n",
    "            \"All lists in data must have the same length as the number of bar labels.\"\n",
    "        )\n",
    "\n",
    "    x = np.arange(len(x_labels))  # the label locations\n",
    "    width = 0.8 / num_groups  # width of each bar\n",
    "\n",
    "    plt.figure(figsize=(10, 6))\n",
    "\n",
    "    for i in range(num_groups):\n",
    "        heights = [v[i] for v in values]\n",
    "        plt.bar(x + i * width, heights, width, label=bar_labels[i])\n",
    "\n",
    "    plt.title(title)\n",
    "    plt.xlabel(xlabel)\n",
    "    plt.ylabel(ylabel)\n",
    "    plt.xticks(x + width * (num_groups - 1) / 2, x_labels)  # Center the tick labels\n",
    "    plt.legend()\n",
    "    plt.grid(axis=\"y\", linestyle=\"--\", alpha=0.6)\n",
    "\n",
    "    if save_path:\n",
    "        plt.savefig(save_path)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab924ba0ac7eabff",
   "metadata": {},
   "source": [
    "## 2. Declare Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be415aef84e070cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = OxfordFlowerDataset(purpose=\"train\")\n",
    "val_dataset = OxfordFlowerDataset(purpose=[\"validation\", \"test\"])\n",
    "\n",
    "train_indices = range(0, len(train_dataset), IMAGE_STEP)\n",
    "val_indices = range(0, len(val_dataset), IMAGE_STEP)\n",
    "print(\"Number of images in the store:\", len(train_indices))\n",
    "print(\"Number of queries:\", len(val_indices))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b63a88e78e223b90",
   "metadata": {},
   "source": [
    "## **3. Load the ClipEmbedder**\n",
    "\n",
    "We will load the embedder on the pre-trained `openai` weights from the `ViT-B/32` variant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "763831fb769f4137",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedder = ClipEmbedder(variant=\"ViT-B/32\", pretrained=\"openai\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "246dd4be523ea126",
   "metadata": {},
   "source": [
    "## **4. Performance metrics**\n",
    "\n",
    "First, we prepare the data. The selected training images are embedded and indexed in an `InMemoryImageEmbeddingStore`, from which the most similar images are retrieved for each query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b41f8223d09379ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_paths = [train_dataset.image_paths[i] for i in train_indices]\n",
    "store = InMemoryImageEmbeddingStore(image_paths=train_paths, embedder=embedder)\n",
    "store.build_store()\n",
    "dataset_labels_dict = dict(\n",
    "    zip(train_dataset.image_paths, train_dataset.labels, strict=True)\n",
    ")\n",
    "val_labels = [val_dataset.labels[i] for i in val_indices]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d6b73d975357b37",
   "metadata": {},
   "source": [
    "### **5.1. Top-k accuracy**\n",
    "\n",
    "How it works:\n",
    "- For each query, retrieve **top-k** most similar images.\n",
    "- If any of them share the same label as the query, that counts as correct.\n",
    "- The final accuracy is `num_correct_queries / num_queries`.\n",
    "\n",
    "Let's compute the top-1 accuracy (the most relevant match has to be the correct one):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23424c86ecbb01c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Top-1 Accuracy for CLIP\n",
    "acc_k1_clip = top_k_accuracy(\n",
    "    images=(val_dataset[i][0] for i in val_indices),\n",
    "    image_labels=val_labels,\n",
    "    store=store,\n",
    "    path_labels_dict=dataset_labels_dict,\n",
    "    k=1,\n",
    ")\n",
    "print(\"Top-1 Accuracy, CLIP:\", acc_k1_clip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aba5dd7bfdb19f1",
   "metadata": {},
   "source": [
    "Normally, we might also consider the second, third and so on.. most relevant results. In this case, we can set `k > 1`. Let's try for `k=5`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "153f3bf9522c1f1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Top-5 Accuracy for CLIP\n",
    "acc_k5_clip = top_k_accuracy(\n",
    "    images=(val_dataset[i][0] for i in val_indices),\n",
    "    image_labels=val_labels,\n",
    "    store=store,\n",
    "    path_labels_dict=dataset_labels_dict,\n",
    "    k=5,\n",
    ")\n",
    "print(\"Top-5 Accuracy, CLIP:\", acc_k5_clip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4dce5e120d66eda",
   "metadata": {},
   "source": [
    "### **5.2. Compute the mAP**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2edbd7d5eebd301",
   "metadata": {},
   "source": [
    "How it works:\n",
    "- If `k` is given, we only consider the `top-k` ranked results per query.\n",
    "- if `k=None` or omitted, we consider all results (the entire dataset).\n",
    "- For each query, we compute average precision (AP). Then we average across all queries, yielding mean average precision (mAP).\n",
    "\n",
    "Example:\n",
    "Image `a` has label `1`, and the top-6 retrieved images have labels:\n",
    "- Truth Labels: [0, 1, 1, 0 ,0, 1]\n",
    "\n",
    "**a) k=None**: we consider all results.\n",
    "- Rank 2: Precision = 1/2\n",
    "- Rank 3: Precision = 2/3\n",
    "- Rank 6: Precision = 3/6\n",
    "- AP = (1/2 + 2/3 + 3/6) / 3 = 0.556\n",
    "\n",
    "**b) k=3**:\n",
    "- Rank 2: Precision = 1/2\n",
    "- Rank 3: Precision = 2/3\n",
    "- AP = (1/2 + 2/3) / 2 = 0.583\n",
    "\n",
    "First, we do it for the whole dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1250e7005751359",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mAP for CLIP\n",
    "mAP_value_clip = top_k_map(\n",
    "    images=(val_dataset[i][0] for i in val_indices),\n",
    "    image_labels=val_labels,\n",
    "    store=store,\n",
    "    path_labels_dict=dataset_labels_dict,\n",
    ")\n",
    "print(\"Mean Average Precision (mAP), CLIP:\", mAP_value_clip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff0af2df9a03ed12",
   "metadata": {},
   "source": [
    "Normally, we might only care about the top results. Let's compute the mAP for the top 5 results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ae67e5962caaa8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mAP of the top-5 results for CLIP\n",
    "mAP_value_top5_clip = top_k_map(\n",
    "    images=(val_dataset[i][0] for i in val_indices),\n",
    "    image_labels=val_labels,\n",
    "    store=store,\n",
    "    path_labels_dict=dataset_labels_dict,\n",
    "    k=5,\n",
    ")\n",
    "print(\"Mean Average Precision (mAP) for Top-5, CLIP:\", mAP_value_top5_clip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10e00fd4d0a04d5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot a bar chart of the mAP and top-k accuracy of the embedder.\n",
    "plot_and_save_barplot(\n",
    "    {\"CLIP\": [mAP_value_clip, acc_k1_clip, acc_k5_clip]},\n",
    "    bar_labels=[\"mAP\", \"Top-1 Accuracy\", \"Top-5 Accuracy\"],\n",
    "    title=\"Performance Metrics for the CLIP Embedder\",\n",
    "    ylabel=\"Value\",\n",
    "    xlabel=\"Embedder\",\n",
    ")"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}