app.py

# /// 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()