image-search-engine

A small, honest semantic image search engine. You embed a set of images into fixed length vectors, drop them into a cosine similarity index, and then query with another image to get back its nearest neighbours. Query with a picture that is already in the collection and it comes back as the top hit. Query with a lightly altered copy of an indexed picture and the original outranks everything unrelated.

The whole thing runs on CPU and the test suite needs no model download.

How it works

There are two pieces.

Encoders turn an image into a vector.

• ImageEncoder is a deterministic encoder built straight from pixels. It concatenates downsampled intensity grids at a few resolutions with per channel colour histograms, then L2 normalizes the result. There is no model to fetch, so it is fast and reproducible, and it is strong enough that near duplicates land close together in cosine space. This is what the tests and the quick demo use.
• TimmEncoder wraps any timm backbone (a small ResNet, a ViT) and pools the final features into one vector. This is the encoder you would point at real photos. You can run it with random weights for an offline experiment or with pretrained=True once you are happy to download weights. If you later want a stronger self supervised backbone, DINOv2 or I-JEPA or a CLIP image tower all drop into the same interface.

The index stores normalized embeddings and ranks them by cosine similarity, which is a plain dot product once everything is unit length. Two backends share one query interface:

• exact computes the full similarity with numpy and always returns the true neighbours. Right for the small collections here.
• approx uses scikit-learn NearestNeighbors with a cosine metric for larger collections.

Quick start

import numpy as np
from src.encoder import ImageEncoder
from src.index import ImageSearchIndex
from src.synthetic import make_collection

# Build a tiny synthetic collection of (n, H, W, 3) images in [0, 1].
images = make_collection(n=12, size=32, seed=0)

encoder = ImageEncoder()
embeddings = encoder.encode(images)

index = ImageSearchIndex(dim=encoder.dim, backend="exact")
index.add(embeddings, ids=[f"img_{i}" for i in range(len(images))])

# Query with an image already in the collection.
hits = index.query_one(embeddings[3], k=3)
for h in hits:
    print(h.item_id, round(h.score, 4))
# img_3 1.0  ... then the two next closest images

Each result is a SearchResult with the integer position, the cosine score, and whatever item_id you attached.

Using a timm backbone instead

from src.encoder import TimmEncoder

encoder = TimmEncoder("resnet18", pretrained=False)  # or pretrained=True
embeddings = encoder.encode(images)

The rest of the pipeline is identical because both encoders return an (n, dim) float32 matrix.

Running the tests

pip install -r requirements.txt
python -m pytest tests/ -q

The tests are behaviour checks rather than fixed numbers. They confirm that embeddings come out unit normalized, that a zero image does not produce NaNs, that querying with an indexed image returns that same image at the top, that a near duplicate beats every unrelated image, that results come back sorted, and that the exact and approximate backends agree on the top hit. On this machine all 27 tests pass in about a second on CPU.

Layout

src/
  encoder.py     ImageEncoder (pixel based) and TimmEncoder (timm backbone)
  index.py       ImageSearchIndex with exact and approx cosine backends
  synthetic.py   gradient and blob image generators plus a near duplicate maker
tests/
  test_encoder.py
  test_index.py

Notes

The synthetic generator makes structured images, smooth colour gradients with a soft blob, rather than pure noise, so two different images are clearly apart and a perturbed copy stays near its source. That is what lets the near duplicate test mean something. No benchmark numbers are baked in beyond what an actual run on your machine reports.