• 【Milvus的以文搜图】

    0. 介绍

    以文搜图指的是，根据文本描述，从图像数据库中检索与文本内容相似的图像数据并返回。通过在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安装

    1. # 通过pip安装
    2. pip install cn_clip

     功能测试

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

    1. import torch
    2. from PIL import Image
    3.  
    4. import cn_clip.clip as clip
    5. from cn_clip.clip import load_from_name, available_models
    6. print("Available models:", available_models())
    7. # Available models: ['ViT-B-16', 'ViT-L-14', 'ViT-L-14-336', 'ViT-H-14', 'RN50']
    8.  
    9. device = "cpu"
    10. if torch.cuda.is_available():
    11.     device = "cuda"
    12.  
    13. model, preprocess = load_from_name("RN50", \
    14.                                     device=device,
    15.                                     download_root="./")
    16.  
    17. model.eval()
    18.  
    19. # image preprocess
    20.  
    21. image_data = Image.open("examples/pokemon.jpeg")
    22. infer_data = preprocess(image_data).unsqueeze(0).to(device)
    23.  
    24. # text data
    25. text_data = clip.tokenize(["杰尼龟", "妙蛙种子", "小火龙", "皮卡丘"]).to(device)
    26.  
    27. with torch.no_grad():
    28.     # image_features = model.encode_image(infer_data)
    29.     # text_features  = model.encode_text(text_data)
    30.  
    31.     # # 对特征进行归一化，请使用归一化后的图文特征用于下游任务
    32.     # image_features /= image_features.norm(dim=-1, keepdim=True) 
    33.     # text_features /= text_features.norm(dim=-1, keepdim=True)
    34.  
    35.     logits_per_image, logits_per_text = model.get_similarity(infer_data, text_data)
    36.     probs = logits_per_image.softmax(dim=-1).cpu().numpy()
    37.  
    38. print("Label probs:", probs)

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

     1.2 提取图文特征

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

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

    准备数据集

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

     文件结构如下所示：

     图文特征提取

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

    1. #!/bin/bash
    2.  
    3. # Usage: extract image and text features
    4. # only supports single-GPU inference
    5. export CUDA_VISIBLE_DEVICES=${1}
    6. export PYTHONPATH=${PYTHONPATH}:`pwd`/cn_clip
    7.  
    8. SPLIT=${2}
    9. DATAPATH=${3}
    10. RESUME=${DATAPATH}/pretrained_weights/clip_cn_rn50.pt
    11. DATASET_NAME=Flickr30k-CN
    12.  
    13. python -u cn_clip/eval/extract_features.py \
    14.     --extract-image-feats \
    15.     --extract-text-feats \
    16.     --image-data="${DATAPATH}/datasets/${DATASET_NAME}/lmdb/${SPLIT}/imgs" \
    17.     --text-data="${DATAPATH}/datasets/${DATASET_NAME}/${SPLIT}_texts.jsonl" \
    18.     --img-batch-size=32 \
    19.     --text-batch-size=32 \
    20.     --context-length=52 \
    21.     --resume=${RESUME} \
    22.     --vision-model=RN50 \
    23.     --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进行检索的结果对比。

    1. #!/bin/bash
    2.  
    3. # Usage: topK search
    4. # only supports single-GPU inference
    5. export CUDA_VISIBLE_DEVICES=${1}
    6. export PYTHONPATH=${PYTHONPATH}:`pwd`/cn_clip
    7.  
    8. SPLIT=${2}
    9. DATAPATH=${3}
    10. RESUME=${DATAPATH}/pretrained_weights/clip_cn_rn50.pt
    11. DATASET_NAME=Flickr30k-CN
    12.  
    13. python -u cn_clip/eval/make_topk_predictions.py \
    14.     --image-feats="${DATAPATH}/datasets/${DATASET_NAME}/${SPLIT}_imgs.img_feat.jsonl" \
    15.     --text-feats="${DATAPATH}/datasets/${DATASET_NAME}/${SPLIT}_texts.txt_feat.jsonl" \
    16.     --top-k=10 \
    17.     --eval-batch-size=32768 \
    18.     --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个内容与文本相似的图像结果。

    2. Milvus

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

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

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

    1. 构建图像特征库

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

    2. 文本检索

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

    2.1 构建特征提取网络

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

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

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

    1. import torch
    2. from PIL import Image
    3. import os
    4. import json
    5. from tqdm import tqdm
    6. import cn_clip.clip as clip
    7. from cn_clip.clip import available_models
    8. print("Available models:", available_models())
    9. # Available models: ['ViT-B-16', 'ViT-L-14', 'ViT-L-14-336', 'ViT-H-14', 'RN50']
    10.  
    11.  
    12.  
    13. class Chinese_CLIP:
    14.     def __init__(self) -> None:
    15.         self.device = "cpu"
    16.         if torch.cuda.is_available():
    17.             self.device = "cuda"
    18.         self.splits = ["valid"]
    19.         self.model, self.preprocess = clip.load_from_name(  "RN50", \
    20.                                                             device=self.device,
    21.                                                             download_root="./")
    22.         self.model.eval()
    23.  
    24.     def extract_image_features(self, img_name):        # for upload
    25.         image_data = Image.open(img_name).convert("RGB")
    26.         infer_data = self.preprocess(image_data)
    27.         infer_data = infer_data.unsqueeze(0).to(self.device)
    28.         with torch.no_grad():
    29.             image_features = self.model.encode_image(infer_data)
    30.         image_features /= image_features.norm(dim=-1, keepdim=True) 
    31.  
    32.         return image_features.cpu().numpy()[0]     # [1, 1024]
    33.  
    34.     def extract_text_features(self, text):            # for search
    35.  
    36.         text_data = clip.tokenize([text]).to(self.device)
    37.         with torch.no_grad():
    38.             text_features = self.model.encode_text(text_data)
    39.         text_features /= text_features.norm(dim=-1, keepdim=True)
    40.         return text_features.cpu().numpy()[0]      # [1, 1024]
    41.     
    42.     def build_Flickr30kCN_image_db(self, data_path):        # for load 
    43.         image_ids = list()
    44.         image_feats = list()
    45.         for split in self.splits:
    46.             image_path = os.path.join(data_path, split + "_imgs.img_feat.jsonl")
    47.             if not os.path.isfile(image_path):
    48.                 print("error, file {} is not exist.".format(image_path))
    49.                 continue
    50.             with open(image_path, "r") as fin:
    51.                 for line in tqdm(fin, desc="build image {} part: ".format(split)):
    52.                     obj = json.loads(line.strip())
    53.                     image_ids.append(obj['image_id'])
    54.                     image_feats.append(obj['feature'])
    55.         return image_ids, image_feats
    56.     
    57.     
    58.     def load_images(self, data_path):        # for mysql
    59.         image_dicts = dict()
    60.         for split in self.splits:
    61.             file_path = os.path.join(data_path, split + "_imgs.tsv")
    62.             if not os.path.isfile(file_path):
    63.                 print("error, file {} is not exist.".format(file_path))
    64.             with open(file_path, "r") as fin_imgs:
    65.                 for line in tqdm(fin_imgs):
    66.                     line = line.strip()
    67.                     image_id, b64 = line.split("\t")
    68.                     image_dicts[image_id] = b64.encode()
    69.         return image_dicts

    2.2 设计collection

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

    1.   def create_collection(self, collection_name):       # 创建collection对象
    2.         # Create milvus collection if not exists
    3.         try:
    4.             if not self.has_collection(collection_name):
    5.                 field1 = FieldSchema(name="id", dtype=DataType.INT64, descrition="int64", is_primary=True, auto_id=True)
    6.                 field2 = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, descrition="float vector",
    7.                                      dim=VECTOR_DIMENSION, is_primary=False)
    8.                 schema = CollectionSchema(fields=[field1, field2], description="collection description")
    9.                 self.collection = Collection(name=collection_name, schema=schema)
    10.                 LOGGER.debug(f"Create Milvus collection: {collection_name}")
    11.             else:
    12.                 # utility.drop_collection(collection_name)
    13.                 # self.create_collection(collection_name)
    14.                 self.set_collection(collection_name)
    15.             return "OK"
    16.         except Exception as e:
    17.             LOGGER.error(f"Failed to create collection to Milvus: {e}")
    18.             sys.exit(1)

    2.3 设计MySQL表

    与Milvus对应的，MySQL table的设计如下，milvus index是milvus生成的索引index，image id是图像对应的图像名，image data是base64编码后的二进制图像数据。

    1.     def create_mysql_table(self, table_name):
    2.         # Create mysql table if not exists
    3.         self.test_connection()
    4.         sql = "create table if not exists " + table_name + "(milvus_id TEXT, image_id TEXT, image_data MEDIUMBLOB not null );"
    5.         try:
    6.             self.cursor.execute(sql)
    7.             LOGGER.debug(f"MYSQL create table: {table_name} with sql: {sql}")
    8.         except Exception as e:
    9.             LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
    10.             sys.exit(1)
    11.     def load_data_to_mysql(self, table_name, data):
    12.         # Batch insert (Milvus_ids, img_path) to mysql
    13.         self.test_connection()
    14.         sql = "insert into " + table_name + " (milvus_id,image_id, image_data) values (%s,%s, %s);"
    15.         try:
    16.             self.cursor.executemany(sql, data)
    17.             self.conn.commit()
    18.             LOGGER.debug(f"MYSQL loads data to table: {table_name} successfully")
    19.         except Exception as e:
    20.             LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
    21.             sys.exit(1)

    2.4 构建图像数据库

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

    1. # Get the vector of images
    2. def extract_features(data_path, model):
    3.     try:
    4.         image_ids, image_feats = model.build_Flickr30kCN_image_db(data_path)
    5.         print(f"Extracting feature from {len(image_ids)} images in total")
    6.             
    7.         return image_ids, image_feats
    8.     except Exception as e:
    9.         LOGGER.error(f"Error with extracting feature from image {e}")
    10.         sys.exit(1)
    11.  
    12.  
    13. # Combine the id of the vector and the name of the image into a list
    14. def format_data(ids, image_ids, image_dicts):
    15.     data = []
    16.     for i in range(len(ids)):
    17.         value = (str(ids[i]), str(image_ids[i]), image_dicts[str(image_ids[i])])
    18.         data.append(value)
    19.     return data
    20.  
    21.  
    22. # Import vectors to Milvus and data to Mysql respectively
    23. def do_load(table_name, image_dir, model, milvus_client, mysql_cli):
    24.     if not table_name:
    25.         table_name = DEFAULT_TABLE
    26.     # 利用模型提取图像特征得到特征向量
    27.     image_ids, vectors = extract_features(image_dir, model)
    28.  
    29.     ids = milvus_client.insert(table_name, vectors)
    30.     milvus_client.create_index(table_name)
    31.  
    32.     image_dicts = model.load_images(image_dir)
    33.  
    34.     mysql_cli.create_mysql_table(table_name)
    35.     mysql_cli.load_data_to_mysql(table_name, format_data(ids, image_ids, image_dicts))
    36.     return len(ids)

    2.5 文本检索

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

    1. def do_search(table_name, text_content, top_k, model, milvus_client, mysql_cli):
    2.     try:
    3.         if not table_name:
    4.             table_name = DEFAULT_TABLE
    5.         features = model.extract_text_features(text_content)
    6.         vectors = milvus_client.search_vectors(table_name, [features], top_k)
    7.         vids = [str(x.id) for x in vectors[0]]
    8.         paths = mysql_cli.search_by_milvus_ids(vids, table_name)
    9.         distances = [x.distance for x in vectors[0]]
    10.         return paths, distances
    11.     except Exception as e:
    12.         LOGGER.error(f"Error with search : {e}")
    13.         sys.exit(1)

    2.6  配置文件

    同样，需要修改config文件中collection名称和特征的维度。

    3. FastApi服务端

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

    1. @app.post('/img/search')
    2. async def search_images( table_name: str = None, content: str = None, topk: int = Form(TOP_K)):
    3.     # Search the upload image in Milvus/MySQL
    4.     try:
    5.         shutil.rmtree(UPLOAD_PATH)
    6.         os.mkdir(UPLOAD_PATH)
    7.         paths, distances = do_search(table_name, content, topk, MODEL, MILVUS_CLI, MYSQL_CLI)
    8.         res = list()
    9.         for data, dist in zip(paths, distances):
    10.             image_path = os.path.join(UPLOAD_PATH, data["image_id"] + ".jpg")
    11.             with open(image_path, "wb") as f:
    12.                 img_data = base64.b64decode(data["image_data"])
    13.                 f.write(img_data)
    14.  
    15.             res.append({
    16.                 "distance": dist,
    17.                 "image_id": image_path,
    18.                 # "image_data": data["image_data"]
    19.             })
    20.         # res = sorted(res.items(), key=lambda item: item[1])
    21.         res.sort(key=lambda item: item["distance"])
    22.         LOGGER.info("Successfully searched similar images!")
    23.         return res
    24.     except Exception as e:
    25.         LOGGER.error(e)
    26.         return {'status': False, 'msg': e}, 400

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

    3.1 构建图像数据库

    进入/img/load条目，输入collection名称和图像特征文件存放的目录，执行。

     后台返回日志如下：

     数据文件目录结果如下：

    3.2 执行文本检索

    本文以Flickr30k-CN数据集中valid子集作为实验，其中的一些文本描述如下：

    1. {"text_id": 148915, "text": "这个小女孩穿着古怪的有条纹的裤子是图，你必须要问，这是更大的，她的脚和她的疲倦的微笑。", "image_ids": [104816788]}
    2. {"text_id": 148916, "text": "一个熟睡的婴儿是在某人的手臂上，穿着一件粉红色条纹的衣服。", "image_ids": [104816788]}
    3. {"text_id": 148917, "text": "一个熟睡的婴儿在一个粉红色的条纹衣服。", "image_ids": [104816788]}
    4. {"text_id": 148918, "text": "在某人的怀里抱着一个小婴儿。", "image_ids": [104816788]}
    5. {"text_id": 148919, "text": "一个小宝宝穿着紫色的裤子。", "image_ids": [104816788]}
    6. {"text_id": 148920, "text": "一个穿着蓝色游泳裤的男孩在水里滑了一个黄色的幻灯片，在水里漂浮着充气玩具。", "image_ids": [1077546505]}
    7. {"text_id": 148921, "text": "一个孩子正从一个漂浮在水面上的彩色气球上掉下来。", "image_ids": [1077546505]}
    8. {"text_id": 148922, "text": "一个男孩在一个池塘里滑了一个彩色的管子。", "image_ids": [1077546505]}
    9. {"text_id": 148923, "text": "一个穿着蓝色短裤的男孩在一个游泳池里滑了一个幻灯片。", "image_ids": [1077546505]}
    10. {"text_id": 148924, "text": "一个男孩骑到一个小后院游泳池。", "image_ids": [1077546505]}
    11. {"text_id": 148925, "text": "两家人正在厨房准备食物。", "image_ids": [1079013716]}
    12. {"text_id": 148926, "text": "厨师为顾客准备汉堡。", "image_ids": [1079013716]}
    13. {"text_id": 148927, "text": "两家厨师在餐厅厨房准备汉堡。", "image_ids": [1079013716]}
    14. {"text_id": 148928, "text": "两家人在厨房准备食物。", "image_ids": [1079013716]}
    15. {"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返回的结果： 

    1. [
    2.   {
    3.     "distance": 0.8329496383666992,
    4.     "image_id": "tmp/search-images/989754491.jpg"
    5.   },
    6.   {
    7.     "distance": 0.8518763780593872,
    8.     "image_id": "tmp/search-images/2508918369.jpg"
    9.   },
    10.   {
    11.     "distance": 0.8582218885421753,
    12.     "image_id": "tmp/search-images/1077546505.jpg"
    13.   },
    14.   {
    15.     "distance": 0.8794834613800049,
    16.     "image_id": "tmp/search-images/154094533.jpg"
    17.   },
    18.   {
    19.     "distance": 0.8805474638938904,
    20.     "image_id": "tmp/search-images/4637950362.jpg"
    21.   },
    22.   {
    23.     "distance": 0.8876401782035828,
    24.     "image_id": "tmp/search-images/1131340021.jpg"
    25.   },
    26.   {
    27.     "distance": 0.8878381848335266,
    28.     "image_id": "tmp/search-images/3621652774.jpg"
    29.   },
    30.   {
    31.     "distance": 0.8895268440246582,
    32.     "image_id": "tmp/search-images/1100214449.jpg"
    33.   },
    34.   {
    35.     "distance": 0.8895866870880127,
    36.     "image_id": "tmp/search-images/811663364.jpg"
    37.   },
    38.   {
    39.     "distance": 0.9175957441329956,
    40.     "image_id": "tmp/search-images/2687539673.jpg"
    41.   }
    42. ]

     通过对比，可以发现两者之间的结果一致。

    4. 总结

    本文，基于中文CLIP和Milvus简单实现以文搜图应用，使用Flickr30k-CN数据集中valid子集作为实验，FastAPI构建可交互的接口，完成小规模的文本检索图像功能。

  • 相关阅读:
    如何监听/抓取两个设备/芯片之间“UART串口”通信数据--监视TXD和RXD
    关于ABAP 中的标题行（又叫表头行或者 header line）的理解
    Python数据分析中常用的工具类和方法【持续更新系列】
    【Mac OS】超好用的几个文字编辑快捷键
    新一轮SocialFi浪潮来袭，Atem Network 再次打响注意力争夺战
    C语言的核心和灵魂
    用 ZEGO Avatar 做一个虚拟人｜虚拟主播直播解决方案
    inscode连接不上gpu，持续8小时，为了数据不丢失续费了6小时，我只想知道什么时候可以连接
    力扣 每日一题 1235. 规划兼职工作【难度：困难，rating: 2022】（动态规划+二分查找）
    【类加载子系统】
  • 原文地址：https://blog.csdn.net/hello_dear_you/article/details/128070563