Chroma 教程

多模态：存图片

🎉摘要：本文介绍如何在 Chroma 中存储图片，实现文本搜图和以图搜图功能。内容包括环境准备、安装依赖、使用 OpenCLIPEmbeddingFunction 和 ImageLoader 编写示例代码，并说明传 uris 与 images 参数的区别及取回图片本体的正确方式。

前面全部是介绍如何向 Chroma 中存储文本信息，但是如果想使用 Chroma 来实现根据文本搜图、或者图搜图，这就需要涉及往 Chroma 存储图片了，下面教你如何做。

环境准备

先准备三张图片，两张狗，一张猫，如下：

运行下面命令安装额外依赖，如果已经安装，则忽略：

pip install pillow open-clip-torch

编写示例代码

Chroma 支持直接存图片（images 参数），但是需要配 OpenCLIPEmbeddingFunction 和 data_loader，如下：

import os
import tempfile
from pathlib import Path
from PIL import Image
from chromadb.utils.data_loaders import ImageLoader
from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction

import chromadb
import numpy as np

# 持久化示例写到临时目录，工作目录保持干净
os.chdir(tempfile.mkdtemp(prefix="chroma_demo_"))

# 基准：基于脚本的图片根目录，不受 os.chdir 影响
image_base = Path(__file__).parent.parent.parent

# 图片完整绝对路径 URI
uri1 = str((image_base / "images" / "cat1.png").resolve())
uri2 = str((image_base / "images" / "dog1.png").resolve())
uri_query = str((image_base / "images" / "dog2.png").resolve())

ef = OpenCLIPEmbeddingFunction()
loader = ImageLoader()

client = chromadb.PersistentClient(path="./chroma_data")
col = client.create_collection(
    "images",
    embedding_function=ef,
    data_loader=loader,
    configuration={"hnsw": {"space": "cosine"}},
)

col.add(
    ids=["img1", "img2"],
    images=[np.array(Image.open(uri1)), np.array(Image.open(uri2))],
    metadatas=[{"label": "cat"}, {"label": "dog"}],
)

# 用文字搜图片
print("用文字搜图片:")
print(col.query(query_texts=["a cute cat"], n_results=2))

# 用图片搜图片
print("\n用图片搜图片:")
print(col.query(query_images=[np.array(Image.open(uri_query))], n_results=2))

# 取回图片本体
print("\n取回图片本体:")
print(col.get(ids=["img1"], include=["data"]))

运行示例，输出如下：

用文字搜图片:
{'ids': [['img1', 'img2']], 'embeddings': None, 'documents': [[None, None]], 'uris': None, 
  'included': ['metadatas', 'documents', 'distances'], 'data': None, 
  'metadatas': [[{'label': 'cat'}, {'label': 'dog'}]], 'distances': [[0.7035367488861084, 0.8313513994216919]]}

用图片搜图片:
{'ids': [['img2', 'img1']], 'embeddings': None, 'documents': [[None, None]], 'uris': None, 
  'included': ['metadatas', 'documents', 'distances'], 'data': None, 
  'metadatas': [[{'label': 'dog'}, {'label': 'cat'}]], 'distances': [[0.3224251866340637, 0.4758944511413574]]}

取回图片本体:
{'ids': ['img1'], 'embeddings': None, 'documents': None, 'uris': [None], 
  'included': ['data', 'uris'], 'data': [None], 'metadatas': [{'label': 'cat'}]}

从上面结果可知，使用“a cute cat”搜索出来的确实是猫的图片，使用 dog2.png 图片搜索出来的也是狗的图片。

但是，取回数据中 data 为 None，没有图片本体。这是因为我们传的是 images= 数组，Chroma 只拿这个数组计算 embedding，不会保存图片路径，data 永远返回 None，ImageLoader 完全不会被调用。

ImageLoader 生效的前提：add 的时候传 uris 参数，不要传 images。Chroma 会把 uri（文件路径）存起来；当你 include=["data"] 读取时，内部调用 ImageLoader，根据 uri 去磁盘重新加载图片，填充 data 字段。例如：

import os
import tempfile
from pathlib import Path
from PIL import Image
from chromadb.utils.data_loaders import ImageLoader
from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction

import chromadb
import numpy as np

# 持久化示例写到临时目录，工作目录保持干净
os.chdir(tempfile.mkdtemp(prefix="chroma_demo_"))

# 基准：基于脚本的图片根目录，不受 os.chdir 影响
image_base = Path(__file__).parent.parent.parent

# 图片完整绝对路径 URI
uri1 = str((image_base / "images" / "cat1.png").resolve())
uri2 = str((image_base / "images" / "dog1.png").resolve())
uri_query = str((image_base / "images" / "dog2.png").resolve())

ef = OpenCLIPEmbeddingFunction()
loader = ImageLoader()

client = chromadb.PersistentClient(path="./chroma_data")
col = client.create_collection(
    "images",
    embedding_function=ef,
    data_loader=loader,
    configuration={"hnsw": {"space": "cosine"}},
)

col.add(
    ids=["img1", "img2"],
    uris=[uri1,uri2], # 看这里
    metadatas=[{"label": "cat"}, {"label": "dog"}],
)

# 用文字搜图片
print("用文字搜图片:")
print(col.query(query_texts=["a cute cat"], n_results=2))

# 用图片搜图片
print("\n用图片搜图片:")
print(col.query(query_images=[np.array(Image.open(uri_query))], n_results=2))

# 取回图片本体
print("\n取回图片本体:")
print(col.get(ids=["img1"], include=["data"]))

运行代码，输出如下：

用文字搜图片:
{'ids': [['img1', 'img2']], 'embeddings': None, 'documents': [[None, None]], 'uris': None,
  'included': ['metadatas', 'documents', 'distances'], 'data': None, 
  'metadatas': [[{'label': 'cat'}, {'label': 'dog'}]], 'distances': [[0.7035367488861084, 0.8313513994216919]]}

用图片搜图片:
{'ids': [['img2', 'img1']], 'embeddings': None, 'documents': [[None, None]], 'uris': None,
  'included': ['metadatas', 'documents', 'distances'], 'data': None, 
  'metadatas': [[{'label': 'dog'}, {'label': 'cat'}]], 'distances': [[0.3224251866340637, 0.4758944511413574]]}

取回图片本体:
{'ids': ['img1'], 'embeddings': None, 'documents': None, 'uris': ['D:\\chroma_db\\images\\cat1.png'], 
'included': ['data', 'uris'], 'data': [array([[[235, 233, 230],
        [236, 235, 231],
        [237, 236, 232],
        ...,
        [200, 191, 186],
        [189, 173, 166],
        [190, 170, 159]]], shape=(200, 200, 3), dtype=uint8)], 'metadatas': [{'label': 'cat'}]}

如果同时设置了 images 和 uris 参数，Chroma 会报错，例如：

col.add(
    ids=["img1", "img2"],
    images=[np.array(Image.open(uri1)), np.array(Image.open(uri2))],
    uris=[uri1, uri2],
    metadatas=[{"label": "cat"}, {"label": "dog"}],
)

错误信息：ValueError: Exactly one of documents, uris, images must be provided in add.

注意：data_loader 也支持从 uri 加载，例如：

col.add(ids=["i1"], uris=["s3://bucket/cat.jpg"], metadatas=[{"label": "cat"}])
res = col.get(ids=["i1"], include=["data"])


上一章：直接用嵌入函数，绕过 Chroma
说说我的看法
全部评论（）
没有评论
关于
本网站专注于 Java、数据库（MySQL、Oracle）、Linux、软件架构及大数据等多领域技术知识分享。涵盖丰富的原创与精选技术文章，助力技术传播与交流。无论是技术新手渴望入门，还是资深开发者寻求进阶，这里都能为您提供深度见解与实用经验，让复杂编码变得轻松易懂，携手共赴技术提升新高度。如有侵权，请来信告知：hxstrive@outlook.com
链接
其他应用
公众号
©2018-2020 hxstrive.com 版权所有 ICP证：蜀ICP备14031461号-2   川公网安备51010802032098