{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "95ab1595",
   "metadata": {},
   "source": [
    "# **Oxford Flowers Clustering Notebook**\n",
    "\n",
    "In this notebook, we:\n",
    "1. Load the validation and test splits of the Oxford Flower Dataset, merging them.\n",
    "2. Compute the CLIP embeddings with the `ClipEmbedder`.\n",
    "3. Cluster these images into 102 clusters (the number of classes).\n",
    "4. Inspect the clusters: the first images of a few clusters and a 3D t-SNE projection of the embeddings.\n",
    "5. Compute RI, ARI, NMI, and interpret the results.\n",
    "6. Repeat the clustering on the similarity matrix instead of the embeddings.\n",
    "\n",
    "---\n",
    "\n",
    "## **1. Setup and Load Data**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "642cc226",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections.abc import Sequence\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "from matplotlib.axes import Axes\n",
    "from sklearn.cluster import SpectralClustering\n",
    "from sklearn.manifold import TSNE\n",
    "from sklearn.metrics import (\n",
    "    adjusted_rand_score,\n",
    "    normalized_mutual_info_score,\n",
    "    rand_score,\n",
    ")\n",
    "\n",
    "from pyvisim.datasets import OxfordFlowerDataset\n",
    "from pyvisim.distance import cosine_similarity\n",
    "from pyvisim.neural_networks import ClipEmbedder\n",
    "from pyvisim.typing import FloatNumpyArray, IntNumpyArray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74a25bc2",
   "metadata": {},
   "source": [
    "### Helpers\n",
    "\n",
    "All helpers used in this notebook are defined here. The first cell clusters the features and scores the result against the ground-truth classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f13e7ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "RANDOM_STATE = 42\n",
    "\n",
    "\n",
    "def cluster_spectrally(features: FloatNumpyArray, n_clusters: int) -> IntNumpyArray:\n",
    "    \"\"\"\n",
    "    Assign every row of ``features`` to one of ``n_clusters`` spectral clusters.\n",
    "\n",
    "    :param features: Array of shape (N, D) with one feature vector per image\n",
    "    :param n_clusters: Number of clusters to form\n",
    "    :return: Cluster id of every image, shape (N,)\n",
    "    \"\"\"\n",
    "    model = SpectralClustering(\n",
    "        n_clusters=n_clusters, affinity=\"nearest_neighbors\", random_state=RANDOM_STATE\n",
    "    )\n",
    "    return np.asarray(model.fit_predict(features))\n",
    "\n",
    "\n",
    "def score_clustering(\n",
    "    true_labels: IntNumpyArray, cluster_labels: IntNumpyArray\n",
    ") -> dict[str, float]:\n",
    "    \"\"\"\n",
    "    Compare a clustering against the ground-truth classes.\n",
    "\n",
    "    :param true_labels: Ground-truth class of every image, shape (N,)\n",
    "    :param cluster_labels: Cluster id of every image, shape (N,)\n",
    "    :return: Rand index, adjusted Rand index and normalized mutual information\n",
    "    \"\"\"\n",
    "    return {\n",
    "        \"RI\": rand_score(true_labels, cluster_labels),\n",
    "        \"ARI\": adjusted_rand_score(true_labels, cluster_labels),\n",
    "        \"NMI\": normalized_mutual_info_score(true_labels, cluster_labels),\n",
    "    }\n",
    "\n",
    "\n",
    "def print_scores(scores: dict[str, float], heading: str) -> None:\n",
    "    \"\"\"\n",
    "    Print clustering scores under a heading.\n",
    "\n",
    "    :param scores: Scores keyed by their name, as returned by :func:`score_clustering`\n",
    "    :param heading: Line printed above the scores\n",
    "    \"\"\"\n",
    "    print(heading)\n",
    "    for name, value in scores.items():\n",
    "        print(f\"{name}: {value:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cf0fac8",
   "metadata": {},
   "source": [
    "The next cell shows the first images that were assigned to a cluster, one row per cluster. Every image is titled with its ground-truth class, so a row with mixed classes reveals an impure cluster at a glance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b65592bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_cluster_samples(\n",
    "    dataset: OxfordFlowerDataset,\n",
    "    cluster_labels: IntNumpyArray,\n",
    "    cluster_ids: Sequence[int],\n",
    "    samples_per_cluster: int,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Show the first images assigned to each cluster, one row per cluster.\n",
    "\n",
    "    :param dataset: Dataset the clustering was computed on\n",
    "    :param cluster_labels: Cluster id of every image in ``dataset``, shape (N,)\n",
    "    :param cluster_ids: Clusters to show, in row order\n",
    "    :param samples_per_cluster: Number of images per row\n",
    "    \"\"\"\n",
    "    _, axes = plt.subplots(\n",
    "        len(cluster_ids),\n",
    "        samples_per_cluster,\n",
    "        figsize=(3 * samples_per_cluster, 3 * len(cluster_ids)),\n",
    "        squeeze=False,\n",
    "    )\n",
    "    for row, cluster_id in zip(axes.tolist(), cluster_ids, strict=True):\n",
    "        members = np.flatnonzero(cluster_labels == cluster_id)[:samples_per_cluster]\n",
    "        draw_cluster_row(row, dataset, members, cluster_id)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def draw_cluster_row(\n",
    "    row: Sequence[Axes],\n",
    "    dataset: OxfordFlowerDataset,\n",
    "    members: IntNumpyArray,\n",
    "    cluster_id: int,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Draw the images of one cluster onto a row of axes.\n",
    "\n",
    "    Axes left over when the cluster has fewer members than the row has axes\n",
    "    are hidden.\n",
    "\n",
    "    :param row: Axes of the row, one per image\n",
    "    :param dataset: Dataset the images are read from\n",
    "    :param members: Indices into ``dataset`` of the images to draw\n",
    "    :param cluster_id: Cluster the images belong to, shown as the row label\n",
    "    \"\"\"\n",
    "    for axis, index in zip(row, members, strict=False):\n",
    "        image, label, _ = dataset[index]\n",
    "        axis.imshow(image)\n",
    "        axis.set_title(f\"Class {label}\")\n",
    "    for axis in row[len(members) :]:\n",
    "        axis.set_visible(False)\n",
    "    for axis in row:\n",
    "        axis.set_xticks([])\n",
    "        axis.set_yticks([])\n",
    "    row[0].set_ylabel(f\"Cluster {cluster_id}\", fontsize=12)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "319df37b",
   "metadata": {},
   "source": [
    "The last two helpers plot data: a 3D t-SNE projection of the embeddings, in which every cluster gets its own colour, and an annotated heatmap of a similarity matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23af99f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def project_to_3d(features: FloatNumpyArray) -> FloatNumpyArray:\n",
    "    \"\"\"\n",
    "    Project feature vectors to three dimensions with t-SNE.\n",
    "\n",
    "    :param features: Array of shape (N, D)\n",
    "    :return: Array of shape (N, 3)\n",
    "    \"\"\"\n",
    "    return TSNE(n_components=3, perplexity=40, random_state=RANDOM_STATE).fit_transform(features)\n",
    "\n",
    "\n",
    "def plot_clusters_3d(\n",
    "    points: FloatNumpyArray,\n",
    "    cluster_labels: IntNumpyArray,\n",
    "    title: str,\n",
    "    colormap_name: str = \"nipy_spectral\",\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Scatter 3D points, colouring the points of a cluster with the same colour.\n",
    "\n",
    "    :param points: Array of shape (N, 3)\n",
    "    :param cluster_labels: Cluster id of every point, shape (N,)\n",
    "    :param title: Title of the plot\n",
    "    :param colormap_name: Colormap the cluster palette is sampled from\n",
    "    \"\"\"\n",
    "    unique_labels, dense_indices = np.unique(cluster_labels, return_inverse=True)\n",
    "    cluster_count = unique_labels.size\n",
    "\n",
    "    # Skip the colormap extremes: they are black/near-white in spectral maps.\n",
    "    palette = plt.get_cmap(colormap_name)(np.linspace(0.05, 0.95, cluster_count))\n",
    "    stride = np.concatenate(\n",
    "        [np.arange(0, cluster_count, 2), np.arange(1, cluster_count, 2)]\n",
    "    )\n",
    "    colours = palette[stride][dense_indices]\n",
    "    colours[cluster_labels < 0] = (0.65, 0.65, 0.65, 0.35)\n",
    "\n",
    "    axis = plt.figure(figsize=(10, 10)).add_subplot(projection=\"3d\")\n",
    "    axis.set_box_aspect(None, zoom=0.85)\n",
    "    axis.scatter(points[:, 0], points[:, 1], points[:, 2], c=colours, s=10)\n",
    "    axis.set_title(title)\n",
    "    axis.set_xlabel(\"t-SNE 1\")\n",
    "    axis.set_ylabel(\"t-SNE 2\")\n",
    "    axis.set_zlabel(\"t-SNE 3\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def show_similarity_heatmap(matrix: FloatNumpyArray, title: str) -> None:\n",
    "    \"\"\"\n",
    "    Draw an annotated heatmap of a similarity matrix.\n",
    "\n",
    "    :param matrix: Square array of pairwise similarities\n",
    "    :param title: Title of the plot\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(8, 7))\n",
    "    sns.heatmap(\n",
    "        matrix,\n",
    "        annot=True,\n",
    "        fmt=\".2f\",\n",
    "        cmap=\"Blues\",\n",
    "        cbar_kws={\"label\": \"Cosine similarity\"},\n",
    "    )\n",
    "    plt.title(title)\n",
    "    plt.xlabel(\"Image index\")\n",
    "    plt.ylabel(\"Image index\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89caf558",
   "metadata": {},
   "source": [
    "### Load the validation and test datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0f808b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = OxfordFlowerDataset(purpose=[\"validation\", \"test\"])\n",
    "true_labels = np.array(dataset.labels)\n",
    "print(\"Number of images in the dataset:\", len(dataset))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f914cf81",
   "metadata": {},
   "source": [
    "## **2. Compute CLIP embeddings**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553017cd",
   "metadata": {},
   "source": [
    "### Define the CLIP embedder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc99f917",
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_embedder = ClipEmbedder(variant=\"ViT-B/32\", pretrained=\"openai\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f94e1d52",
   "metadata": {},
   "source": [
    "### Compute the CLIP embeddings for both the validation and test splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "367bb4b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_embeddings = clip_embedder.embed(image for image, *_ in dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8df62418",
   "metadata": {},
   "source": [
    "## **3. Cluster into 102 Clusters**\n",
    "\n",
    "`102` is the number of classes in the Oxford Flowers dataset. We want to see how well the clustering algorithm can cluster the images into these classes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "388d9f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLASSES = 102\n",
    "cluster_labels = cluster_spectrally(clip_embeddings, n_clusters=NUM_CLASSES)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e83fa32",
   "metadata": {},
   "source": [
    "### Inspect the clusters\n",
    "\n",
    "Before looking at any score, let's look at the clusters themselves. The grid below shows the first 5 images of each of the first 5 clusters. Ideally, all images in a row share the same class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a70adf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_cluster_samples(\n",
    "    dataset, cluster_labels, cluster_ids=range(5), samples_per_cluster=5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "174ef29d",
   "metadata": {},
   "source": [
    "### Visualize the embeddings in 3D\n",
    "\n",
    "t-SNE is used to project the 512-dimensional CLIP embeddings onto three dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98a31f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "projected_embeddings = project_to_3d(clip_embeddings)\n",
    "plot_clusters_3d(\n",
    "    projected_embeddings,\n",
    "    cluster_labels,\n",
    "    title=f\"t-SNE projection of the CLIP embeddings, {NUM_CLASSES} spectral clusters\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02ce10cc",
   "metadata": {},
   "source": [
    "### Compute RI, ARI, NMI\n",
    "\n",
    "The Rand index (RI) is the fraction of image pairs on which the clustering and the ground truth agree. The adjusted Rand index (ARI) corrects it for chance agreement, and the normalized mutual information (NMI) measures how much knowing the cluster of an image tells us about its class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10523d63",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = score_clustering(true_labels, cluster_labels)\n",
    "heading = f\"Spectral clustering of the embeddings into {NUM_CLASSES} clusters:\"\n",
    "print_scores(scores, heading)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d27d6d62",
   "metadata": {},
   "source": [
    "## **4. Cluster on the Similarity Matrix**\n",
    "\n",
    "Now, instead of using the embeddings themselves, we will use the `similarity matrix` of the embeddings to cluster the images. So each row in this matrix will represent the similarity of an image to all other images in the dataset (hence, all diagonal elements will be 1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dbe1b99",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_matrix = cosine_similarity(clip_embeddings, clip_embeddings)\n",
    "show_similarity_heatmap(\n",
    "    similarity_matrix[:10, :10],\n",
    "    title=\"Similarity matrix of the dataset, first 10 images\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e5723a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_cluster_labels = cluster_spectrally(\n",
    "    similarity_matrix, n_clusters=NUM_CLASSES\n",
    ")\n",
    "similarity_scores = score_clustering(true_labels, similarity_cluster_labels)\n",
    "heading = f\"Spectral clustering of the similarity matrix into {NUM_CLASSES} clusters:\"\n",
    "print_scores(similarity_scores, heading)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9db05a1",
   "metadata": {},
   "source": [
    "## **5. Conclusion**\n",
    "\n",
    "We've demonstrated:\n",
    "- How to cluster images directly on CLIP embeddings.\n",
    "- How to inspect the clusters visually, through sample images and a 3D t-SNE projection.\n",
    "- How to compute RI, ARI and NMI for objective evaluation.\n",
    "\n",
    "For CLIP, clustering on the embeddings themselves performs significantly better than clustering on the similarity matrix."
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}