visual-rag

Retrieval augmented generation over images. You give the system a corpus of images, it embeds and indexes them, and then for a natural language query it retrieves the images that answer the query, packs them into a grounded context, and hands that to a generator which writes an answer that cites the images it used.

The point of the repository is the architecture. Every stage is a clean seam you can replace: the embedder, the index, the context builder, and the generator are all independent. The default embedder and generator are chosen so the whole thing runs on CPU with no model downloads and no API keys, which makes the behavior easy to test, but the shape of the code is exactly what you would deploy on top of a learned encoder like CLIP and a hosted language model.

How it works

There are four stages.

Embedding. Images and text queries are projected into one shared vector space so a query vector lands near the images that answer it. The embedding in src/embed.py is built from two real signals. The first is a visual signal read straight from pixels: a per channel color histogram plus brightness, contrast, edge energy, and saturation statistics. The second is a semantic bridge: image tags and query words are hashed into a shared concept subspace using the hashing trick, so a query that mentions "cat" reinforces the same dimensions as an image tagged "cat". Text queries have no pixels, so their visual block is zero and they retrieve through the concept subspace. Images carry both blocks. Everything is L2 normalized so cosine similarity is well defined.

This is a genuine, deterministic encoder rather than a placeholder. Because it lives behind the ImageEmbedder interface, swapping in CLIP later means rewriting one file and nothing else.

Indexing. src/index.py stores the normalized embeddings as a matrix and scores a query with a single matrix vector product, which is how a flat vector index works in practice. Each vector travels with its image id, tags, and caption so retrieval returns enough to ground a generator. You can query by text or by a probe image.

Context assembly. src/context.py turns the ranked retrieval results into an ordered, citable context. Each retrieved image becomes a compact evidence block (id, score, tags, caption). This is the visual analogue of concatenating retrieved passages into a prompt, and it is where you would cap the prompt budget by limiting the number of blocks.

Generation. src/generate.py defines the generator seam. The default StubGenerator is deterministic and offline: it reads the assembled context and writes an answer that cites the retrieved image ids in rank order. It can only cite evidence that is present in the context, which is the core grounding contract a real RAG generator must honor. To use an actual language model you implement the Generator protocol and return a GroundedAnswer; the rest of the pipeline does not change.

src/pipeline.py wires the four stages into VisualRAG, which takes a corpus in and returns a grounded answer out.

Usage

import numpy as np
from src import VisualRAG

rag = VisualRAG()
rag.add_image(
    "img_cat",
    np.clip(np.ones((32, 32, 3)) * [0.85, 0.6, 0.2], 0, 1),
    tags=["cat", "animal", "indoor"],
    caption="An orange tabby cat sitting indoors.",
)
rag.add_image(
    "img_ocean",
    np.clip(np.ones((32, 32, 3)) * [0.1, 0.3, 0.7], 0, 1),
    tags=["ocean", "water", "blue"],
    caption="A blue ocean horizon.",
)
rag.build()

out = rag.answer("an orange cat indoors", top_k=2)
print(out.results[0].image_id)   # img_cat
print(out.answer.text)           # an answer citing [img_cat]
print(out.answer.citations)      # ('img_cat', ...)

out carries the ranked retrieval results, the assembled context, and the grounded answer, so you can inspect every stage.

Swapping in a real model

The two replaceable parts are the embedder and the generator.

For embeddings, subclass or replace ImageEmbedder so embed_image and embed_text return vectors from a learned encoder such as CLIP, then hand the new embedder to VisualRAG. The index, context, and generation code is unchanged because it only depends on the vector interface.

For generation, implement the Generator protocol with a call to a hosted or local language model. Feed it context.render(), parse its reply into a GroundedAnswer, and keep the rule that it may only cite ids that appear in the context.

Tests

The tests are behavior checks, not snapshots. They confirm that retrieval returns the relevant image first for a query, that queries for different concepts retrieve different images, that the assembled context contains every retrieved piece of evidence in rank order, and that the generator cites the retrieved images and never cites anything outside the context. There is also a full pipeline test that runs retrieve, assemble, and generate on a cat query and checks the answer is grounded in the cat image.

Run them with:

python -m pytest tests/ -q

The synthetic corpus is generated in memory, so the suite needs no network and no API key. On this machine all 24 tests pass in well under a second.

Layout

src/
  embed.py      shared image and text embedding space
  index.py      cosine similarity vector index over the corpus
  context.py    assemble retrieved images into a grounded context
  generate.py   pluggable generator, deterministic stub by default
  pipeline.py   VisualRAG end to end
tests/
  conftest.py   synthetic image corpus fixtures
  test_embed.py
  test_retrieval.py
  test_context_and_generation.py