Milvus的以文搜图

842次阅读
没有评论
  1. 介绍
    以文搜图指的是,根据文本描述,从图像数据库中检索与文本内容相似的图像数据并返回。通过在CSDN中搜索以文搜图,找到了如下两篇文章:

从零到一,教你搭建「以文搜图」搜索服务(一)Zilliz Planet的博客-CSDN博客以文搜图 搜索引擎技术
从零到一,教你搭建「CLIP 以文搜图」搜索服务(二):5 分钟实现原型_Zilliz Planet的博客-CSDN博客
这两篇文章的内容介绍了以文搜图的原理,以及如何使用Faiss来搭建以文搜图的应用。

与上述两篇文章不同,本文主要讲述如何基于Chinese-CLIP和Milvus来搭建一个以文搜图应用。

  1. Chinese-CLIP
    CLIP是OpenAI在2021年提出的多模态神经网络模型,该模型基于OpenAI 收集到的 4 亿对图像文本对进行训练,分别将文本和图像进行编码,之后使用 metric learning 进行权重,其目标是将图像与文本的相似性提高,大致如下图所示。具体内容不在本文赘述。 Chinese-CLIP是CLIP模型的中文版本,使用大规模的图文对进行训练得到,针对中文领域数据能够实现更好的效果。本文的以文搜图应用使用中文CLIP作为文本和图像的特征提取器,提取数据集中所有图像的特征作为数据库,用于检索。

1.1 安装中文CLIP
中文CLIP的README文件详细描述了如何安装、跨模态检索的Finetune、预测和评估,以及零样本分类等内容,本文主要基于该文件,进行相关实验。

依赖环境
安装中文CLIP库之前,需要满足如下环境配置:

python >= 3.6.4
pytorch >= 1.8.0 (with torchvision >= 0.9.0)
CUDA Version >= 10.2
pip安装

通过pip安装

pip install cn_clip
功能测试
通过下面的代码,可以验证中文CLIP库是否安装正确,主要内容是提取给定的图像和文本特征,并计算它们之间的相似度。

import torch
from PIL import Image

import cn_clip.clip as clip
from cn_clip.clip import load_from_name, available_models
print(“Available models:”, available_models())

Available models: [‘ViT-B-16’, ‘ViT-L-14’, ‘ViT-L-14-336’, ‘ViT-H-14’, ‘RN50’]

device = “cpu”
if torch.cuda.is_available():
device = “cuda”

model, preprocess = load_from_name(“RN50”, \
device=device,
download_root=”./”)

model.eval()

image preprocess

image_data = Image.open(“examples/pokemon.jpeg”)
infer_data = preprocess(image_data).unsqueeze(0).to(device)

text data

text_data = clip.tokenize([“杰尼龟”, “妙蛙种子”, “小火龙”, “皮卡丘”]).to(device)

with torch.no_grad():
# image_features = model.encode_image(infer_data)
# text_features = model.encode_text(text_data)

# # 对特征进行归一化,请使用归一化后的图文特征用于下游任务
# image_features /= image_features.norm(dim=-1, keepdim=True) 
# text_features /= text_features.norm(dim=-1, keepdim=True)

logits_per_image, logits_per_text = model.get_similarity(infer_data, text_data)
probs = logits_per_image.softmax(dim=-1).cpu().numpy()

print(“Label probs:”, probs)
现阶段,中文CLIP提供5种不同规模的模型,如下图所示,上述例子使用了RN50,probs输出pookemon图像与4个文本之间的相似性。

1.2 提取图文特征
在README文件中,跨模态检索部分介绍了如何在下游数据集中对模型进行finetune、预测和评估模型性能。

本节中,主要内容是基于中文clip如何提取图文特征,即使用现有的预训练模型,不做finetune,对数据集进行推理,得到数据集图文特征文件。

准备数据集
按照介绍,为了提高效率,需要将数据集转换为tsv和jsonl格式的文件,然后将其转换为内存索引的LMDB数据库文件。为了简化实验的步骤,本文直接使用项目中提供的已经转换好的数据集,进行后续的实验,如下图所示,下载文件放到对应文件目录下。本文以Flickr30K-CN数据集进行实验。

文件结构如下所示:

图文特征提取
通过执行如下命令,能够完成对数据集中图像和文本的特征提取。

!/bin/bash

Usage: extract image and text features

only supports single-GPU inference

export CUDA_VISIBLE_DEVICES=${1}
export PYTHONPATH=${PYTHONPATH}:pwd/cn_clip

SPLIT=${2}
DATAPATH=${3}
RESUME=${DATAPATH}/pretrained_weights/clip_cn_rn50.pt
DATASET_NAME=Flickr30k-CN

python -u cn_clip/eval/extract_features.py \
–extract-image-feats \
–extract-text-feats \
–image-data=”${DATAPATH}/datasets/${DATASET_NAME}/lmdb/${SPLIT}/imgs” \
–text-data=”${DATAPATH}/datasets/${DATASET_NAME}/${SPLIT}_texts.jsonl” \
–img-batch-size=32 \
–text-batch-size=32 \
–context-length=52 \
–resume=${RESUME} \
–vision-model=RN50 \
–text-model=RBT3-chinese
其中第一个参数是用于推理的GPU ID,第二参数DATAPATH指的数据存放的目录,split指的是train、test或者valid。图片特征保存于${split}_imgs.img_feat.jsonl文件,文本特征则保存于${split}_texts.txt_feat.jsonl。

KNN检索
在中文CLIP项目中,提供了一个简单的KNN检索,用来计算文到图、图到文检索的top-k召回结果,本文主要计算从文到图的检索结果,与后续使用Milvus进行检索的结果对比。

!/bin/bash

Usage: topK search

only supports single-GPU inference

export CUDA_VISIBLE_DEVICES=${1}
export PYTHONPATH=${PYTHONPATH}:pwd/cn_clip

SPLIT=${2}
DATAPATH=${3}
RESUME=${DATAPATH}/pretrained_weights/clip_cn_rn50.pt
DATASET_NAME=Flickr30k-CN

python -u cn_clip/eval/make_topk_predictions.py \
–image-feats=”${DATAPATH}/datasets/${DATASET_NAME}/${SPLIT}_imgs.img_feat.jsonl” \
–text-feats=”${DATAPATH}/datasets/${DATASET_NAME}/${SPLIT}_texts.txt_feat.jsonl” \
–top-k=10 \
–eval-batch-size=32768 \
–output=”${DATAPATH}/datasets/${DATASET_NAME}/${SPLIT}_predictions.jsonl”
结果保存到${split}_predictions.jsonl文件中,每行表示一个文本召回的top-k图片id,格式如下:

{“text_id”: 153915, “image_ids”: [5791244, 1009692167, 7454547004, 3564007203, 38130571, 2525270674, 2195419145, 2503091968, 4966265765, 3690431163]}
至此,对于中文CLIP如何提取图像和文本特征有了大致的了解,接下来,将中文CLIP提取的图像特征存储到Milvus中,然后输入文本,CLIP提取文本特征,在图像数据库中检索得到topK个内容与文本相似的图像结果。

  1. Milvus
    如前面文章介绍,Milvus是一个向量数据库,基于深度学习网络提取的特征进行对象之间相似度计算,返回topK个相似的结果。

不同的检索任务,需要对milvus和mysql的数据内容进行设计,以及如何提取对象的特征是十分重要的。

以文搜图的整体流程如上图所示,图片来自链接:

  1. 构建图像特征库

图像数据集通过多模态模型提取特征,本文中使用中文CLIP多模态模型,然后将特征存储到Milvus中,用于后续的检索。

  1. 文本检索

请求文本通过多模态模型提取特征,得到文本特征,然后从Milvus数据库中进行检索,得到topK个相似的内容。

2.1 构建特征提取网络
如下方的代码所示,基于中文CLIP构建一个特征提取类,使用extract_image_featuresh和extract_text_features提取图像和文本特征。

build_Flickr30kCN_image_db函数用于从第一部分生成的图像特征文件中提取图像ID和特征来构建图像数据库,用于后续的检索任务。

load_images函数用于提取图像ID和以base64编码的二进制图像数据,后续插入到MySQL中,检索时,返回相应的图像内容。

import torch
from PIL import Image
import os
import json
from tqdm import tqdm
import cn_clip.clip as clip
from cn_clip.clip import available_models
print(“Available models:”, available_models())

Available models: [‘ViT-B-16’, ‘ViT-L-14’, ‘ViT-L-14-336’, ‘ViT-H-14’, ‘RN50’]

class Chinese_CLIP:
def init(self) -> None:
self.device = “cpu”
if torch.cuda.is_available():
self.device = “cuda”
self.splits = [“valid”]
self.model, self.preprocess = clip.load_from_name( “RN50”, \
device=self.device,
download_root=”./”)
self.model.eval()

def extract_image_features(self, img_name):        # for upload
    image_data = Image.open(img_name).convert("RGB")
    infer_data = self.preprocess(image_data)
    infer_data = infer_data.unsqueeze(0).to(self.device)
    with torch.no_grad():
        image_features = self.model.encode_image(infer_data)
    image_features /= image_features.norm(dim=-1, keepdim=True) 

    return image_features.cpu().numpy()[0]     # [1, 1024]

def extract_text_features(self, text):            # for search

    text_data = clip.tokenize([text]).to(self.device)
    with torch.no_grad():
        text_features = self.model.encode_text(text_data)
    text_features /= text_features.norm(dim=-1, keepdim=True)
    return text_features.cpu().numpy()[0]      # [1, 1024]

def build_Flickr30kCN_image_db(self, data_path):        # for load 
    image_ids = list()
    image_feats = list()
    for split in self.splits:
        image_path = os.path.join(data_path, split + "_imgs.img_feat.jsonl")
        if not os.path.isfile(image_path):
            print("error, file {} is not exist.".format(image_path))
            continue
        with open(image_path, "r") as fin:
            for line in tqdm(fin, desc="build image {} part: ".format(split)):
                obj = json.loads(line.strip())
                image_ids.append(obj['image_id'])
                image_feats.append(obj['feature'])
    return image_ids, image_feats


def load_images(self, data_path):        # for mysql
    image_dicts = dict()
    for split in self.splits:
        file_path = os.path.join(data_path, split + "_imgs.tsv")
        if not os.path.isfile(file_path):
            print("error, file {} is not exist.".format(file_path))
        with open(file_path, "r") as fin_imgs:
            for line in tqdm(fin_imgs):
                line = line.strip()
                image_id, b64 = line.split("\t")
                image_dicts[image_id] = b64.encode()
    return image_dicts               

2.2 设计collection
和其他任务一样,以文搜图任务的collection内容设计如下,milvus index是milvus生成的索引index,embedding是图像的特征。

def create_collection(self, collection_name): # 创建collection对象
# Create milvus collection if not exists
try:
if not self.has_collection(collection_name):
field1 = FieldSchema(name=”id”, dtype=DataType.INT64, descrition=”int64″, is_primary=True, auto_id=True)
field2 = FieldSchema(name=”embedding”, dtype=DataType.FLOAT_VECTOR, descrition=”float vector”,
dim=VECTOR_DIMENSION, is_primary=False)
schema = CollectionSchema(fields=[field1, field2], description=”collection description”)
self.collection = Collection(name=collection_name, schema=schema)
LOGGER.debug(f”Create Milvus collection: {collection_name}”)
else:
# utility.drop_collection(collection_name)
# self.create_collection(collection_name)
self.set_collection(collection_name)
return “OK”
except Exception as e:
LOGGER.error(f”Failed to create collection to Milvus: {e}”)
sys.exit(1)
2.3 设计MySQL表
与Milvus对应的,MySQL table的设计如下,milvus index是milvus生成的索引index,image id是图像对应的图像名,image data是base64编码后的二进制图像数据。

def create_mysql_table(self, table_name):
    # Create mysql table if not exists
    self.test_connection()
    sql = "create table if not exists " + table_name + "(milvus_id TEXT, image_id TEXT, image_data MEDIUMBLOB not null );"
    try:
        self.cursor.execute(sql)
        LOGGER.debug(f"MYSQL create table: {table_name} with sql: {sql}")
    except Exception as e:
        LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
        sys.exit(1)
def load_data_to_mysql(self, table_name, data):
    # Batch insert (Milvus_ids, img_path) to mysql
    self.test_connection()
    sql = "insert into " + table_name + " (milvus_id,image_id, image_data) values (%s,%s, %s);"
    try:
        self.cursor.executemany(sql, data)
        self.conn.commit()
        LOGGER.debug(f"MYSQL loads data to table: {table_name} successfully")
    except Exception as e:
        LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
        sys.exit(1)

2.4 构建图像数据库
在operations目录下的load文件的功能是构建图像特征库,本文是从第一部分的文件中解析得到图像ID和图像特征,并将其插入到milvus中,将图像id和图像数据插入MySQL中。

Get the vector of images

def extract_features(data_path, model):
try:
image_ids, image_feats = model.build_Flickr30kCN_image_db(data_path)
print(f”Extracting feature from {len(image_ids)} images in total”)

    return image_ids, image_feats
except Exception as e:
    LOGGER.error(f"Error with extracting feature from image {e}")
    sys.exit(1)

Combine the id of the vector and the name of the image into a list

def format_data(ids, image_ids, image_dicts):
data = []
for i in range(len(ids)):
value = (str(ids[i]), str(image_ids[i]), image_dicts[str(image_ids[i])])
data.append(value)
return data

Import vectors to Milvus and data to Mysql respectively

def do_load(table_name, image_dir, model, milvus_client, mysql_cli):
if not table_name:
table_name = DEFAULT_TABLE
# 利用模型提取图像特征得到特征向量
image_ids, vectors = extract_features(image_dir, model)

ids = milvus_client.insert(table_name, vectors)
milvus_client.create_index(table_name)

image_dicts = model.load_images(image_dir)

mysql_cli.create_mysql_table(table_name)
mysql_cli.load_data_to_mysql(table_name, format_data(ids, image_ids, image_dicts))
return len(ids)

2.5 文本检索
在operations目录下的search文件功能是根据给定的文本内容,执行图像检索,具体实现如下:

def do_search(table_name, text_content, top_k, model, milvus_client, mysql_cli):
try:
if not table_name:
table_name = DEFAULT_TABLE
features = model.extract_text_features(text_content)
vectors = milvus_client.search_vectors(table_name, [features], top_k)
vids = [str(x.id) for x in vectors[0]]
paths = mysql_cli.search_by_milvus_ids(vids, table_name)
distances = [x.distance for x in vectors[0]]
return paths, distances
except Exception as e:
LOGGER.error(f”Error with search : {e}”)
sys.exit(1)
2.6 配置文件
同样,需要修改config文件中collection名称和特征的维度。

  1. FastApi服务端
    在main文件中,由于任务发生了变化,因此需要对search条目的函数进行修改,输入文本内容,将输出的base64编码的图像数据解码保存下来。

@app.post(‘/img/search’)
async def search_images( table_name: str = None, content: str = None, topk: int = Form(TOP_K)):
# Search the upload image in Milvus/MySQL
try:
shutil.rmtree(UPLOAD_PATH)
os.mkdir(UPLOAD_PATH)
paths, distances = do_search(table_name, content, topk, MODEL, MILVUS_CLI, MYSQL_CLI)
res = list()
for data, dist in zip(paths, distances):
image_path = os.path.join(UPLOAD_PATH, data[“image_id”] + “.jpg”)
with open(image_path, “wb”) as f:
img_data = base64.b64decode(data[“image_data”])
f.write(img_data)

        res.append({
            "distance": dist,
            "image_id": image_path,
            # "image_data": data["image_data"]
        })
    # res = sorted(res.items(), key=lambda item: item[1])
    res.sort(key=lambda item: item["distance"])
    LOGGER.info("Successfully searched similar images!")
    return res
except Exception as e:
    LOGGER.error(e)
    return {'status': False, 'msg': e}, 400

同样,使用uvicorn main:app –reload命令启动fastapi服务,在浏览器中进入docs页面。

3.1 构建图像数据库
进入/img/load条目,输入collection名称和图像特征文件存放的目录,执行。

后台返回日志如下:

数据文件目录结果如下:

3.2 执行文本检索
本文以Flickr30k-CN数据集中valid子集作为实验,其中的一些文本描述如下:

{“text_id”: 148915, “text”: “这个小女孩穿着古怪的有条纹的裤子是图,你必须要问,这是更大的,她的脚和她的疲倦的微笑。”, “image_ids”: [104816788]}
{“text_id”: 148916, “text”: “一个熟睡的婴儿是在某人的手臂上,穿着一件粉红色条纹的衣服。”, “image_ids”: [104816788]}
{“text_id”: 148917, “text”: “一个熟睡的婴儿在一个粉红色的条纹衣服。”, “image_ids”: [104816788]}
{“text_id”: 148918, “text”: “在某人的怀里抱着一个小婴儿。”, “image_ids”: [104816788]}
{“text_id”: 148919, “text”: “一个小宝宝穿着紫色的裤子。”, “image_ids”: [104816788]}
{“text_id”: 148920, “text”: “一个穿着蓝色游泳裤的男孩在水里滑了一个黄色的幻灯片,在水里漂浮着充气玩具。”, “image_ids”: [1077546505]}
{“text_id”: 148921, “text”: “一个孩子正从一个漂浮在水面上的彩色气球上掉下来。”, “image_ids”: [1077546505]}
{“text_id”: 148922, “text”: “一个男孩在一个池塘里滑了一个彩色的管子。”, “image_ids”: [1077546505]}
{“text_id”: 148923, “text”: “一个穿着蓝色短裤的男孩在一个游泳池里滑了一个幻灯片。”, “image_ids”: [1077546505]}
{“text_id”: 148924, “text”: “一个男孩骑到一个小后院游泳池。”, “image_ids”: [1077546505]}
{“text_id”: 148925, “text”: “两家人正在厨准备食物。”, “image_ids”: [1079013716]}
{“text_id”: 148926, “text”: “厨师为顾客准备汉堡。”, “image_ids”: [1079013716]}
{“text_id”: 148927, “text”: “两家厨师在餐厅厨房准备汉堡。”, “image_ids”: [1079013716]}
{“text_id”: 148928, “text”: “两家人在厨房准备食物。”, “image_ids”: [1079013716]}
{“text_id”: 148929, “text”: “厨房烹饪中的2个厨师。”, “image_ids”: [1079013716]}
在search条目中,使用如上的一个描述进行检索

返回的topK结果如下:

3.3 结果对比
通过对比milvus返回的结果和中文CLIP项目中KNN检索的结果,来判断检索的正确性。找到上述文本对应的text ID,在valid_predictions.jsonl文件中,找到对应的topK个Image ID。

KNN 检索结果:

{“text_id”: 148924, “image_ids”: [989754491, 2508918369, 1077546505, 154094533, 4637950362, 1131340021, 3621652774, 1100214449, 811663364, 2687539673]}
Milvus返回的结果:

[
{
“distance”: 0.8329496383666992,
“image_id”: “tmp/search-images/989754491.jpg”
},
{
“distance”: 0.8518763780593872,
“image_id”: “tmp/search-images/2508918369.jpg”
},
{
“distance”: 0.8582218885421753,
“image_id”: “tmp/search-images/1077546505.jpg”
},
{
“distance”: 0.8794834613800049,
“image_id”: “tmp/search-images/154094533.jpg”
},
{
“distance”: 0.8805474638938904,
“image_id”: “tmp/search-images/4637950362.jpg”
},
{
“distance”: 0.8876401782035828,
“image_id”: “tmp/search-images/1131340021.jpg”
},
{
“distance”: 0.8878381848335266,
“image_id”: “tmp/search-images/3621652774.jpg”
},
{
“distance”: 0.8895268440246582,
“image_id”: “tmp/search-images/1100214449.jpg”
},
{
“distance”: 0.8895866870880127,
“image_id”: “tmp/search-images/811663364.jpg”
},
{
“distance”: 0.9175957441329956,
“image_id”: “tmp/search-images/2687539673.jpg”
}
]
通过对比,可以发现两者之间的结果一致。

  1. 总结
    本文,基于中文CLIP和Milvus简单实现以文搜图应用,使用Flickr30k-CN数据集中valid子集作为实验,FastAPI构建可交互的接口,完成小规模的文本检索图像功能。
正文完
可以使用微信扫码关注公众号(ID:xzluomor)
post-qrcode
 
评论(没有评论)