# /// script
# requires-python = ">=3.12"
# dependencies = [
# "chromadb==1.0.4",
# "datasets==3.5.0",
# "marimo",
# "matplotlib==3.10.1",
# "numpy==2.2.4",
# "open-clip-torch==2.32.0",
# "pillow==11.1.0",
# ]
# ///
import marimo
__generated_with = "0.12.8"
app = marimo.App(width="medium")
@app.cell
def _():
import marimo as mo
return (mo,)
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
# Multimodal Retrieval
Chroma supports multimodal collections, i.e. collections which contain, and can be queried by, multiple modalities of data.
This notebook shows an example of how to create and query a collection with both text and images, using Chroma's built-in features.
"""
)
return
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
## Dataset
We us a small subset of the [coco object detection dataset](https://huggingface.co/datasets/detection-datasets/coco), hosted on HuggingFace.
We download a small fraction of all the images in the dataset locally, and use it to create a multimodal collection.
"""
)
return
@app.cell
def _():
import os
from datasets import load_dataset
from matplotlib import pyplot as plt
return load_dataset, os
@app.cell
def _(load_dataset, mo):
with mo.status.spinner(title="Loading dataset"):
dataset = load_dataset(
path="detection-datasets/coco",
name="default",
split="train",
streaming=True,
)
N_IMAGES = 20
return N_IMAGES, dataset
@app.cell
def _(N_IMAGES, dataset, mo, os):
# Write the images to a folder
IMAGE_FOLDER = "images"
os.makedirs(IMAGE_FOLDER, exist_ok=True)
i = 0
all_images = []
with mo.status.spinner(title="Loading images"):
for row in dataset.take(N_IMAGES):
image = row["image"]
all_images.append(image)
image.save(f"images/{i}.jpg")
i += 1
return IMAGE_FOLDER, all_images
@app.cell(hide_code=True)
def _(mo):
img_width = mo.ui.slider(
label="Image width", start=100, stop=300, step=10, debounce=True
)
img_width
return (img_width,)
@app.cell(hide_code=True)
def _(all_images, img_width, mo):
import io
def as_image(src):
img_byte_arr = io.BytesIO()
src.save(img_byte_arr, format=src.format or "PNG")
img_byte_arr.seek(0)
return mo.image(img_byte_arr, width=img_width.value)
mo.hstack(
[as_image(_img) for _img in all_images[10:]],
wrap=True,
)
return
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
## Ingesting multimodal data
Chroma supports multimodal collections by referencing external URIs for data types other than text.
All you have to do is specify a data loader when creating the collection, and then provide the URI for each entry.
For this example, we are only adding images, though you can also add text.
"""
)
return
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
### Creating a multi-modal collection
First we create the default Chroma client.
"""
)
return
@app.cell
def _():
import chromadb
client = chromadb.Client()
return (client,)
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
Next we specify an embedding function and a data loader.
The built-in `OpenCLIPEmbeddingFunction` works with both text and image data. The `ImageLoader` is a simple data loader that loads images from a local directory.
"""
)
return
@app.cell
def _():
from chromadb.utils.data_loaders import ImageLoader
from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
embedding_function = OpenCLIPEmbeddingFunction()
image_loader = ImageLoader()
return embedding_function, image_loader
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""We create a collection with the embedding function and data loader.""")
return
@app.cell
def _(IMAGE_FOLDER, client, embedding_function, image_loader, os):
collection = client.create_collection(
name="multimodal_collection",
embedding_function=embedding_function,
data_loader=image_loader,
get_or_create=True,
)
# Get the uris to the images
image_uris = sorted(
[
os.path.join(IMAGE_FOLDER, image_name)
for image_name in os.listdir(IMAGE_FOLDER)
]
)
ids = [str(i) for i in range(len(image_uris))]
collection.add(ids=ids, uris=image_uris)
return (collection,)
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
### Adding multi-modal data
We add image data to the collection using the image URIs. The data loader and embedding functions we specified earlier will ingest data from the provided URIs automatically.
"""
)
return
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
## Querying a multi-modal collection
We can query the collection using text as normal, since the `OpenCLIPEmbeddingFunction` works with both text and images.
"""
)
return
@app.cell(hide_code=True)
def _(mo):
query = mo.ui.text_area(label="Query with text", full_width=True).form(
bordered=False
)
mo.vstack([query, mo.md("Try: *animal* or *vehicle*")])
return (query,)
@app.cell
def _(collection, mo, query):
mo.stop(not query.value)
_retrieved = collection.query(
query_texts=[query.value], include=["data"], n_results=3
)
[mo.image(img, height=200) for img in _retrieved["data"][0]]
return
@app.cell(hide_code=True)
def _(mo):
mo.md(
r"""
/// admonition | One more thing!
We can also query by images directly, by using the `query_images` field in the `collection.query` method.
///
"""
)
return
@app.cell
def _(collection, mo, selected_image):
mo.stop(not selected_image.value)
import numpy as np
from PIL import Image
query_image = np.array(Image.open(selected_image.path()))
selected = mo.as_html(mo.image(query_image))
_retrieved = collection.query(
query_images=[query_image], include=["data"], n_results=5
)
results = [mo.image(_img) for _img in _retrieved["data"][0][1:]]
return results, selected
@app.cell(hide_code=True)
def _(IMAGE_FOLDER, mo):
selected_image = mo.ui.file_browser(IMAGE_FOLDER, multiple=False)
selected_image
return (selected_image,)
@app.cell(hide_code=True)
def _(mo, results, selected):
mo.hstack(
[
mo.vstack([mo.md("## Selected"), selected]),
mo.vstack([mo.md("## Similar"), *results]),
],
widths="equal",
gap=4,
)
return
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""This example was adapted from [multimodal_retrieval.ipynb](https://github.com/chroma-core/chroma/blob/main/examples/multimodal/multimodal_retrieval.ipynb), using `marimo convert`.""")
return
if __name__ == "__main__":
app.run()