{ "cells": [ { "cell_type": "markdown", "id": "508f1760", "metadata": {}, "source": [ "# Image Search\n", "\n", "This notebook builds a searchable image gallery and queries it with an image rather than with text. The gallery is the `Oxford 102 Flowers` dataset [1], each image is turned into a feature vector by a pre-trained `CLIP` model [2], and the vectors are indexed in an `HNSW` graph [3] so that the nearest neighbors of a query are found without scanning the gallery.\n", "\n", "The pipeline has four steps:\n", "\n", "1. **Embed**: every gallery image is read and passed through `ClipEmbedder`, which returns one vector per image.\n", "2. **Index**: the vectors are handed to the search index that `InMemoryImageEmbeddingStore` owns from then on.\n", "3. **Serialization**: the store is written to a single `safetensors` file, embedder included, so the gallery never has to be embedded twice.\n", "4. **Query**: an unseen image is embedded with the same embedder and matched against the index.\n", "\n", "> [!NOTE]\n", ">\n", "> Note that embedding the whole gallery takes a while. Once the file is on disk, the notebook can be restarted from the loading cell.\n", "\n", "## `Hnsw` Search Index\n", "\n", "A brute-force search compares the query against every gallery vector, which is exact but linear in the size of the gallery. `HNSW` instead builds a multi-layer proximity graph over the vectors and walks it greedily, which reaches the neighborhood of the query in far fewer comparisons. The trade is recall: the answer is approximate, and how approximate is controlled by the parameters set below.\n", "\n", "Further reading: [`Image Store`](https://mechacritter.github.io/Python-Visual-Similarity/image_similarity_retrieval/image_store/in_memory_image_embedding_store/in_memory_image_embedding_store.html) and [`External Indexes`](https://mechacritter.github.io/Python-Visual-Similarity/image_similarity_retrieval/image_store/external_search_index/external_search_index.html)." ] }, { "cell_type": "markdown", "id": "cc445be9", "metadata": {}, "source": [ "## Import libraries" ] }, { "cell_type": "code", "execution_count": null, "id": "a847f200", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "from PIL import Image\n", "\n", "from pyvisim.datasets import OxfordFlowerDataset\n", "from pyvisim.neural_networks import ClipEmbedder\n", "from pyvisim.retrieval.image_store import InMemoryImageEmbeddingStore" ] }, { "cell_type": "markdown", "id": "92509656", "metadata": {}, "source": [ "## Hyperparameters\n", "\n", "> [!NOTE]\n", ">\n", "> `IMAGE_STEP` keeps every `IMAGE_STEP`-th image of the training split in the gallery. Set to `1` if you want to use all images." ] }, { "cell_type": "code", "execution_count": null, "id": "ec789a32", "metadata": {}, "outputs": [], "source": [ "IMAGE_STEP = 4" ] }, { "cell_type": "markdown", "id": "ce2bbe57", "metadata": {}, "source": [ "## Build the image store\n", "\n", "Constructing `InMemoryImageEmbeddingStore` only records the gallery. `build_store()` reads every path in `image_paths`, embeds it with `embedder` and hands the resulting matrix to the search index, so call it right after the constructor. Until then the store holds nothing, and searching it raises `RuntimeError`. Pass `lazy_build=False` to have the constructor call `build_store()` itself.\n", "\n", "The gallery is every `IMAGE_STEP`-th image of the training split of `Oxford 102 Flowers` [1], which holds 6149 images.\n", "\n", "Arguments:\n", "\n", "- `embedder`: any object exposing an `embed` method. `ClipEmbedder` is used here, but the `VLAD`, `Fisher Vector`, `Siamese` and `Triplet` embedders fit here just as well.\n", "- `search_index`: `\"hnsw\"` builds the graph, `None` falls back to an exact brute-force scan. An `ExternalSearchIndex` can be passed instead to search through a `FAISS` index [5].\n", "- `index_params`: forwarded to the index constructor. `graph_degree` is the number of bidirectional links created per node of the graph, `build_candidates` the size of the candidate list kept while building it. Both raise recall, `graph_degree` at the cost of memory and `build_candidates` at the cost of build time.\n", "\n", "The metric space defaults to `\"cosine\"`, which stores the vectors L2-normalised and scores them by `1 - cosine_similarity`." ] }, { "cell_type": "code", "execution_count": null, "id": "585f96c2", "metadata": {}, "outputs": [], "source": [ "train_dataset = OxfordFlowerDataset()\n", "train_indices = range(0, len(train_dataset), IMAGE_STEP)\n", "train_image_paths = [train_dataset.image_paths[i] for i in train_indices]\n", "print(\"Number of images in the gallery:\", len(train_image_paths))\n", "\n", "embedder = ClipEmbedder()\n", "\n", "image_store = InMemoryImageEmbeddingStore(\n", " image_paths=train_image_paths,\n", " embedder=embedder,\n", " search_index=\"hnsw\",\n", " index_params={\"graph_degree\": 16, \"build_candidates\": 200},\n", ")\n", "image_store.build_store()\n", "image_store.save_to_disk(\"flower_image_store.safetensors\")" ] }, { "cell_type": "markdown", "id": "222d8794", "metadata": {}, "source": [ "## Load the store back from disk\n", "\n", "The saved file carries the embedder as well, so the store is rebuilt without touching the original images and without downloading the `CLIP` weights again. The index is rebuilt from the saved embeddings, using the parameters it was saved with.\n", "\n", "This is the entry point for a fresh session. Run the import cell, then this one, and skip the embedding cell above entirely. A store loaded from disk is already built, so it needs no `build_store()` call.\n", "\n", "One caveat applies to external indexes only. A `FAISS` index cannot be written into the file, so a rebuilt one has to be handed back as `search_index=...`. Without it the store falls back to an exact brute-force scan over the saved embeddings and warns about the fallback." ] }, { "cell_type": "code", "execution_count": null, "id": "3a6b6c74", "metadata": {}, "outputs": [], "source": [ "image_store = InMemoryImageEmbeddingStore.load_from_disk(\n", " \"flower_image_store.safetensors\"\n", ")" ] }, { "cell_type": "markdown", "id": "6b96b063", "metadata": {}, "source": [ "## Helper to display images>" ] }, { "cell_type": "code", "execution_count": null, "id": "cf588ef5", "metadata": {}, "outputs": [], "source": [ "def _gallery_labels_by_path():\n", " dataset = OxfordFlowerDataset()\n", " return dict(zip(dataset.image_paths, dataset.labels, strict=True))\n", "\n", "\n", "#: Class id of every gallery image, keyed by its path.\n", "GALLERY_LABELS_BY_PATH = _gallery_labels_by_path()\n", "\n", "\n", "def _as_rgb_array(image_array_or_path):\n", " if isinstance(image_array_or_path, np.ndarray):\n", " return image_array_or_path\n", " with Image.open(image_array_or_path) as image:\n", " return np.asarray(image.convert(\"RGB\"))\n", "\n", "\n", "def visualize_image(image_array_or_path, title=None):\n", " plt.figure(figsize=(4, 4))\n", " plt.imshow(_as_rgb_array(image_array_or_path))\n", " plt.axis(\"off\")\n", " if title is not None:\n", " plt.title(title)\n", " plt.show()\n", "\n", "\n", "def visualize_candidate(candidate):\n", " label = GALLERY_LABELS_BY_PATH[candidate.path]\n", " visualize_image(\n", " candidate.path, title=f\"label {label}, score = {candidate.score:.4f}\"\n", " )" ] }, { "cell_type": "markdown", "id": "1d7d5cb3", "metadata": {}, "source": [ "## The query image\n", "\n", "The query is drawn from the test split, so it was never embedded into the gallery and the search has to generalise to an unseen image. Its own class label is shown above it, and it is the label every retrieved image is measured against further down." ] }, { "cell_type": "code", "execution_count": null, "id": "e75e3eb9", "metadata": {}, "outputs": [], "source": [ "test_dataset = OxfordFlowerDataset(purpose=\"test\")\n", "image1, label1, _ = test_dataset[172]\n", "visualize_image(image1, title=f\"query, label {label1}\")" ] }, { "cell_type": "markdown", "id": "97ca7fa1", "metadata": {}, "source": [ "## Retrieve the nearest neighbors\n", "\n", "`retrieve_top_k_similar` embeds the query with the store's own embedder, walks the index, and maps the returned row ids back to gallery paths.\n", "\n", "It takes a batch as readily as a single image and always returns one ranked list per query, which is why the first list is unpacked with `[0]`. Batching is worth using whenever there is more than one query, because the index answers a single `(M, D)` matrix far faster than it answers `M` separate queries.\n", "\n", "Each result is a `Candidate`, a named tuple of `path` and `score`, ordered best first. In cosine space the score is `1 - cosine_similarity`, so **lower means more similar** and a near-duplicate of the query scores close to `0.0`.\n", "\n", "Query time is bounded by `search_candidates`, the width of the graph walk, which defaults to `50`. Asking for more than `search_candidates` neighbors raises it to `k` automatically." ] }, { "cell_type": "code", "execution_count": null, "id": "d48abb83", "metadata": {}, "outputs": [], "source": [ "candidates = image_store.retrieve_top_k_similar(image1, k=5)[0]" ] }, { "cell_type": "markdown", "id": "b5e54fe8", "metadata": {}, "source": [ "## Inspect the results\n", "\n", "The top matches should carry the query's own class label. `HNSW` is approximate, so a graph built for speed can miss a true neighbor. If the results look poor, rebuild with a larger `graph_degree` or `build_candidates`, or raise `search_candidates` at query time. Building the store with `search_index=None` gives the exact ranking to compare against, at the cost of scanning the whole gallery on every query." ] }, { "cell_type": "code", "execution_count": null, "id": "58d9bfb3", "metadata": {}, "outputs": [], "source": [ "for candidate in candidates:\n", " visualize_candidate(candidate)" ] }, { "cell_type": "markdown", "id": "61c056f8", "metadata": {}, "source": [ "## Refine the query with alpha query expansion\n", "\n", "A query is a single point, and a single point can sit off to the side of the group it belongs to. Alpha query expansion (aQE) [7] moves it back towards the middle of that group before the ranking is decided.\n", "\n", "The store searches once, reads the embeddings of the `expansion_neighbors` best matches back off the index, and replaces the query by the L2-normalised weighted average of itself and those matches. Each match is weighted by its cosine similarity to the query raised to `expansion_alpha`, so a close match pulls harder than a distant one, and a match whose similarity is not positive does not pull at all. The final search then runs with the refined query.\n", "\n", "Arguments:\n", "\n", "- `query_expansion`: turns the refinement on. It is off by default, because it costs one extra index search per query plus the decoding of `expansion_neighbors` gallery vectors.\n", "- `expansion_alpha`: exponent of the similarity weights. `0` weights every match with a positive similarity alike, which is the classic average query expansion.\n", "- `expansion_neighbors`: how many of the top-ranked gallery images are averaged into the query.\n", "\n", "The scores below are still the store's own distances, so lower still means more similar, but they are measured from the refined query rather than from the one that was embedded and are not comparable to the scores above." ] }, { "cell_type": "code", "execution_count": null, "id": "77e5461a", "metadata": {}, "outputs": [], "source": [ "expanded_candidates = image_store.retrieve_top_k_similar(\n", " image1,\n", " k=5,\n", " query_expansion=True,\n", " expansion_alpha=3.0,\n", " expansion_neighbors=50,\n", ")[0]\n", "\n", "for expanded_candidate in expanded_candidates:\n", " visualize_candidate(expanded_candidate)" ] }, { "cell_type": "markdown", "id": "710f5482", "metadata": {}, "source": [ "## Re-rank the candidates with k-reciprocal encoding\n", "\n", "Query expansion changes the query. Re-ranking leaves the query untouched and re-orders a pool of candidates that has already been retrieved.\n", "\n", "`KReciprocalReranker` [8] asks of every candidate whether the query lies among its own nearest neighbors, not only whether it lies among the query's. That relation is far stricter than plain proximity: a false match can sit close to the query while the query sits nowhere near the false match's own neighbors. The query and every candidate are encoded into a k-reciprocal feature, a vector that holds a weight for each of their k-reciprocal neighbors and zero elsewhere, and the Jaccard distance between two such features says how much the two neighborhoods agree.\n", "\n", "The neighborhoods are built among the candidates themselves, so a pool no larger than the answer leaves them nothing to say. A pool of `100` is retrieved below and the best five are kept.\n", "\n", "Arguments:\n", "\n", "- `k1`: size of the neighborhoods the k-reciprocal sets are built from. The pool should hold at least this many candidates.\n", "- `k2`: size of the neighborhood the local query expansion averages the features over. `1` turns that step off.\n", "- `lambda_value`: weight of the original distance in the final one, from `0` (Jaccard distance only) to `1` (original ranking kept).\n", "- `top_k`: how many of the re-ranked candidates `rerank` returns.\n", "\n", "The scores are the final distances of the re-ranking rather than the store's. They lie in `[0, 1]`, and lower again means more similar.\n", "\n", "Further reading: [`Re-ranking`](https://mechacritter.github.io/Python-Visual-Similarity/image_similarity_retrieval/reranking/k_reciprocal_reranker/k_reciprocal_reranker.html)." ] }, { "cell_type": "code", "execution_count": null, "id": "dec273c8", "metadata": {}, "outputs": [], "source": [ "from pyvisim.retrieval.reranking import KReciprocalReranker\n", "\n", "reranker = KReciprocalReranker(image_store, k1=20, k2=6, lambda_value=0.3)\n", "\n", "candidate_pool = image_store.retrieve_top_k_similar(image1, k=100)[0]\n", "reranked_candidates = reranker.rerank(candidate_pool, top_k=5)\n", "\n", "for reranked_candidate in reranked_candidates:\n", " visualize_candidate(reranked_candidate)" ] }, { "cell_type": "markdown", "id": "7f0a41c2", "metadata": {}, "source": [ "## Search through a `FAISS` index\n", "\n", "`ExternalSearchIndex` lets the store search through an index built by another library, so an algorithm this package does not ship can still be used over the same gallery. The example below rebuilds the `HNSW` graph with `FAISS` [5] and hands it over.\n", "\n", "The graph is built over `image_store.embeddings`, the matrix the store's own index holds, so the gallery is never embedded a second time.\n", "\n", "Two things are the caller's job here rather than the store's:\n", "\n", "- **The metric.** `METRIC_INNER_PRODUCT` ranks by cosine similarity only if the vectors are L2-normalised. They already are, because the store was built in the default `\"cosine\"` space, and the queries arrive normalised too, because `ClipEmbedder` normalises what it returns.\n", "- **The parameters.** `index_params` is not forwarded to an external index. `m` is a constructor argument of `IndexHNSWFlat`, while `efConstruction` and `efSearch` are set on the graph itself.\n", "\n", "An external index already holds the gallery, so a store on one is built as soon as it is constructed and needs no `build_store()` call.\n", "\n", "`from_faiss_index` reads the gallery vectors back off the index, which `IndexHNSWFlat` can do because it keeps them uncompressed. A quantized index cannot, and needs them passed explicitly as the second argument." ] }, { "cell_type": "code", "execution_count": null, "id": "3c9d5ae1", "metadata": {}, "outputs": [], "source": [ "import faiss\n", "\n", "from pyvisim.retrieval.image_store import ExternalSearchIndex\n", "\n", "gallery_vectors = image_store.embeddings\n", "\n", "faiss_index = faiss.IndexHNSWFlat(\n", " gallery_vectors.shape[1], 16, faiss.METRIC_INNER_PRODUCT\n", ")\n", "faiss_index.hnsw.efConstruction = 200\n", "faiss_index.add(gallery_vectors)\n", "faiss_index.hnsw.efSearch = 50\n", "\n", "faiss_store = InMemoryImageEmbeddingStore(\n", " image_paths=image_store.paths,\n", " embedder=embedder,\n", " search_index=ExternalSearchIndex.from_faiss_index(faiss_index, name=\"faiss-hnsw\"),\n", ")\n", "faiss_store" ] }, { "cell_type": "markdown", "id": "b2ef6d40", "metadata": {}, "source": [ "## Compare the two rankings\n", "\n", "The two indexes report their scores in opposite units. `FAISS` was asked for inner products, so its score is a cosine similarity and **higher means more similar**. The built-in index returns the `1 - cosine_similarity` distance of its `\"cosine\"` space, where lower means more similar. On L2-normalised vectors these are the same quantity read from opposite ends, so `1 - score` converts either one into the other exactly.\n", "\n", "**A conversion is therefore applied below.** The `FAISS` similarity is turned into a cosine distance with `1 - faiss_candidate.score` before it is printed, so both columns speak the same units and can be read against each other directly. Neither store rescales anything of its own accord. The conversion is done here in the cell, and what the index itself returned is still `faiss_candidate.score`.\n", "\n", "Both graphs are approximate and are built by different code, so the two rankings are free to disagree on a borderline neighbor. On this query they do not." ] }, { "cell_type": "code", "execution_count": null, "id": "48c1b9fa", "metadata": {}, "outputs": [], "source": [ "faiss_candidates = faiss_store.retrieve_top_k_similar(image1, k=5)[0]\n", "\n", "for faiss_candidate, hnsw_candidate in zip(faiss_candidates, candidates, strict=True):\n", " agreement = \"same\" if faiss_candidate.path == hnsw_candidate.path else \"differs\"\n", " faiss_distance = 1.0 - faiss_candidate.score\n", " print(\n", " f\"{agreement:>7} faiss={faiss_distance:.4f} \"\n", " f\"hnsw={hnsw_candidate.score:.4f} {faiss_candidate.path}\"\n", " )" ] }, { "cell_type": "markdown", "id": "0d6b8e57", "metadata": {}, "source": [ "## Serialize a store built on an external index\n", "\n", "`save_to_disk` writes the embeddings, the paths and the embedder, but not the `FAISS` index itself, which this library cannot serialize. Loading therefore takes a rebuilt index back as `search_index=...`, and its name is checked against the saved one, so a mismatch is reported. Left out, the store falls back to an exact brute-force scan over the saved embeddings and warns about it.\n", "\n", "A quantized index needs `embeddings=...` passed to `save_to_disk` as well, because what it reconstructs is an approximation of the embeddings rather than the embeddings themselves.\n", "\n", "Note that the file written below is a second full copy of the gallery, roughly the size of the one written further up." ] }, { "cell_type": "code", "execution_count": null, "id": "e4a37b19", "metadata": {}, "outputs": [], "source": [ "faiss_store.save_to_disk(\"flower_faiss_store.safetensors\")\n", "\n", "restored_store = InMemoryImageEmbeddingStore.load_from_disk(\n", " \"flower_faiss_store.safetensors\",\n", " search_index=ExternalSearchIndex.from_faiss_index(faiss_index, name=\"faiss-hnsw\"),\n", ")\n", "restored_candidates = restored_store.retrieve_top_k_similar(image1, k=5)[0]\n", "\n", "for restored_candidate in restored_candidates:\n", " print(f\"{restored_candidate.score:.4f} {restored_candidate.path}\")" ] }, { "cell_type": "markdown", "id": "eb54c0b8", "metadata": {}, "source": [ "## Visualise the results\n", "\n", "It can be observed that the results are identical to the ones returned by the built-in index, since the same `hnsw` parameters are used. \n", "\n", "The scores in the titles are the ones `FAISS` itself returned rather than the distances converted further up, so here **higher means more similar**. The labels are read the same way as everywhere else.\n", "\n", "Now, you can plug in any `FAISS` index you like for the search." ] }, { "cell_type": "code", "execution_count": null, "id": "acb1aced", "metadata": {}, "outputs": [], "source": [ "for faiss_candidate in faiss_candidates:\n", " visualize_candidate(faiss_candidate)" ] }, { "cell_type": "markdown", "id": "795fb47b", "metadata": {}, "source": [ "## References\n", "\n", "[1] Nilsback, M.-E., & Zisserman, A. (2008). Automated Flower Classification\n", "over a Large Number of Classes. In Proceedings of the Sixth Indian Conference\n", "on Computer Vision, Graphics and Image Processing (ICVGIP), 722-729.\n", "https://www.robots.ox.ac.uk/~vgg/data/flowers/102/\n", "\n", "\n", "[2] Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S.,\n", "Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., & Sutskever, I.\n", "(2021). Learning Transferable Visual Models From Natural Language Supervision.\n", "In Proceedings of the 38th International Conference on Machine Learning (ICML),\n", "PMLR 139, 8748-8763. https://arxiv.org/abs/2103.00020\n", "\n", "\n", "[3] Malkov, Y. A., & Yashunin, D. A. (2020). Efficient and Robust Approximate\n", "Nearest Neighbor Search Using Hierarchical Navigable Small World Graphs. IEEE\n", "Transactions on Pattern Analysis and Machine Intelligence, 42(4), 824-836.\n", "https://arxiv.org/abs/1603.09320\n", "\n", "\n", "[4] hnswlib, the reference implementation the compiled index is built on.\n", "https://github.com/nmslib/hnswlib\n", "\n", "\n", "[5] Johnson, J., Douze, M., & Jegou, H. (2019). Billion-Scale Similarity Search\n", "with GPUs. IEEE Transactions on Big Data, 7(3), 535-547.\n", "https://arxiv.org/abs/1702.08734\n", "\n", "\n", "[6] Pinecone. Hierarchical Navigable Small Worlds (HNSW).\n", "https://www.pinecone.io/learn/series/faiss/hnsw/\n", "\n", "\n", "[7] Radenovic, F., Tolias, G., & Chum, O. (2019). Fine-tuning CNN Image\n", "Retrieval with No Human Annotation. IEEE Transactions on Pattern Analysis and\n", "Machine Intelligence, 41(7), 1655-1668. https://arxiv.org/abs/1711.02512\n", "\n", "\n", "[8] Zhong, Z., Zheng, L., Cao, D., & Li, S. (2017). Re-ranking Person\n", "Re-identification with k-reciprocal Encoding. In Proceedings of the IEEE\n", "Conference on Computer Vision and Pattern Recognition (CVPR), 1318-1327.\n", "https://arxiv.org/abs/1701.08398" ] } ], "metadata": {}, "nbformat": 4, "nbformat_minor": 5 }