基于文本嵌入和 CLIP 图像嵌入的多模态检索

!mkdir data
!mkdir images
# 将网页下载成html文件并存储在本地
!curl https://abcbirds.org/blog/backyard-birds/ --output ./data/backyard-birds.html
from pathlib import Path
from llama_index import download_loader
from llama_index import SimpleDirectoryReader
# 将下载的文件作为非结构化文档加载
UnstructuredReader = download_loader('UnstructuredReader')
# 将 ./data/ 目录中的所有非结构化数据文件加载到列表中。
dir_reader = SimpleDirectoryReader('./data', file_extractor={
  ".html": UnstructuredReader()
})
documents = dir_reader.load_data()
print(f'Loaded {len(documents)} documents')

import os
# 存储图像元数据，包括文件名和路径
image_metadata_dict = {}
# 获取 images 文件夹中所有图像文件的列表
image_files = os.listdir("images")
# 循环遍历每个图像文件并添加元数据
for image_file in image_files:
  image_path = os.path.join("images", image_file)
  # 获取图像文件名
  image_filename = os.path.basename(image_path)
  # 获取图片文件路径
  image_file_path = os.path.abspath(image_path)
  # 创建字典来存储图像元数据
  image_metadata_dict[image_filename] = {
    "filename": image_filename,
    "img_path": image_file_path,
  }

from PIL import Image
import matplotlib.pyplot as plt
import os
# 在网格中绘制多个图像的简单而有用的方法
def plot_images(image_metadata_dict):
    images = []
    images_shown = 0
    for image_filename in image_metadata_dict:
        img_path = image_metadata_dict[image_filename]["img_path"]
        if os.path.isfile(img_path):
            # 打开图像文件并将其转换为 RGB 色彩空间。
            filename = image_metadata_dict[image_filename]["filename"]
            image = Image.open(img_path).convert("RGB")
            # 在 8x8 网格的子图中绘制图像，还禁用轴上的刻度标签以使绘图更清晰。
            plt.subplot(8, 8, len(images) + 1)
            plt.imshow(image)
            plt.xticks([])
            plt.yticks([])
            images.append(filename)
            images_shown += 1
            if images_shown >= 64:
                break
    plt.tight_layout()
plot_images(image_metadata_dict)

import qdrant_client
from llama_index import (
    ServiceContext,
    SimpleDirectoryReader,
)
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index import VectorStoreIndex, StorageContext
from llama_index.llms import OpenAI
# create QdrantClient with the location set to ":memory:", which means the vector db will be stored in memory
text_client = qdrant_client.QdrantClient(location=":memory:")
# load text documents from the data directory
text_documents = SimpleDirectoryReader("./data").load_data()
# define llm and embedding model in ServiceContext
service_context = ServiceContext.from_defaults(
    llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
    embed_model="local:BAAI/bge-base-en-v1.5"
)
# create QdrantVectorStore using QdrantClient and the collection name "birds_text"
text_vector_store = QdrantVectorStore(
    client=text_client, collection_name="birds_text"
)
# create StorageContext object using the QdrantVectorStore
storage_context = StorageContext.from_defaults(vector_store=text_vector_store)
# create VectorStoreIndex using the text documents and StorageContext
birds_text_index = VectorStoreIndex.from_documents(
    text_documents,
    storage_context=storage_context,
    service_context=service_context
)
# define text query engine
text_query_engine = birds_text_index.as_query_engine()

import torch
import clip
import numpy as np
# load the CLIP model with the name ViT-B/32
model, preprocess = clip.load("ViT-B/32")
# the resolution of the input images expected by the model
input_resolution = model.visual.input_resolution
# the maximum length of the input text
context_length = model.context_length
# the size of the vocabulary used by the model
vocab_size = model.vocab_size
# print the information about the model to the console
print(
    "Model parameters:",
    f"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}",
)
print("Input resolution:", input_resolution)
print("Context length:", context_length)
print("Vocab size:", vocab_size)

Model parameters: 151,277,313 
Input resolution: 224 
Context length: 77 
Vocab size: 49408

# 设置用于 CLIP 模型的设备，CUDA (GPU) 或 CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
# 存储每个图像的图像嵌入
img_emb_dict = {}
with torch.no_grad():
    # 迭代图像元数据字典并提取每个图像的图像嵌入
    for image_filename in image_metadata_dict:
        img_file_path = image_metadata_dict[image_filename]["img_path"]
        if os.path.isfile(img_file_path):
            image = (
                # 使用 CLIP 模型的预处理功能对图像进行预处理
                # 解压缩图像张量以添加批量维度
                # 将图像张量移动到第 1 行指定的设备
                preprocess(Image.open(img_file_path)).unsqueeze(0).to(device)
            )
            # 使用CLIP模型的 encode_image 函数提取图像特征
            image_features = model.encode_image(image)
            # 将图像特征存储在图像嵌入字典中
            img_emb_dict[image_filename] = image_features

from llama_index.schema import ImageDocument
# 创建 ImageDocument 对象列表，每个对象对应数据集中的每个图像
img_documents = []
for image_filename in image_metadata_dict:
    # img_emb_dict 字典包含图像嵌入
    if image_filename in img_emb_dict:
        filename = image_metadata_dict[image_filename]["filename"]
        filepath = image_metadata_dict[image_filename]["img_path"]
        print(filepath)
        # 为每个图像创建一个 ImageDocument 
        newImgDoc = ImageDocument(
            text=filename, metadata={"filepath": filepath}
        )
        # 在 ImageDocument 上设置图像嵌入
        newImgDoc.embedding = img_emb_dict[image_filename].tolist()[0]
        img_documents.append(newImgDoc)
# 创建 QdrantVectorStore，集合名称为 birds_img 
image_vector_store = QdrantVectorStore(
    client=text_client, collection_name="birds_img"
)
# 定义存储上下文
storage_context = StorageContext.from_defaults(vector_store=image_vector_store)
# 定义图像索引
image_index = VectorStoreIndex.from_documents(
    img_documents,
    storage_context=storage_context
)

from llama_index.vector_stores import VectorStoreQuery
# 将文本查询作为输入并从向量存储中返回最相似的图像
def retrieve_results_from_image_index(query):
    # 首先标记文本查询并将其转换为张量
    text = clip.tokenize(query).to(device)
    # 使用 CLIP 模型对文本张量进行编码以生成查询嵌入
    query_embedding = model.encode_text(text).tolist()[0]
    # 创建一个 VectorStoreQuery
    image_vector_store_query = VectorStoreQuery(
        query_embedding=query_embedding,
        similarity_top_k=1, 
        mode="default",
    )
    # 针对图像矢量存储执行查询
    image_retrieval_results = image_vector_store.query(
        image_vector_store_query
    )
    return image_retrieval_results
# 获取图像检索结果列表并创建一个尺寸为 16 x 5 英寸的新图形
def plot_image_retrieve_results(image_retrieval_results):
    plt.figure(figsize=(16, 5))
    img_cnt = 0
    # 迭代图像检索结果，对于每个结果，在子图中显示相应的图像及其分数。
    # 子图为 2 x 3，图像显示在 img_cnt + 第 1 个子图中。
    # 子图的标题是图像的分数，格式为小数点后四位。
    for returned_image, score in zip(
        image_retrieval_results.nodes, image_retrieval_results.similarities
    ):
        img_name = returned_image.text
        img_path = returned_image.metadata["filepath"]
        image = Image.open(img_path).convert("RGB")
        plt.subplot(2, 3, img_cnt + 1)
        plt.title("{:.4f}".format(score))
        plt.imshow(image)
        plt.xticks([])
        plt.yticks([])
        img_cnt += 1
# 定义 image_query 函数
def image_query(query):
    image_retrieval_results = retrieve_results_from_image_index(query)
    plot_image_retrieve_results(image_retrieval_results)