全部產品
Search
文件中心

Tair:基於Tair Vector實現圖文多模態檢索

更新時間:Jun 30, 2024

本文介紹基於Tair Vector和CLIP實現即時高效能圖文多模態檢索的解決方案。

背景資訊

在互連網中,大量資訊(例如圖片、文本等)通常以非結構化的形式存在。達摩院的CLIP開源模型內建了Text transformer、ResNet等模型,支援對圖片、文本等非結構化資料進行特徵提取,並解析成結構化資料。

您可以先通過CLIP模型將圖片、文檔等資料預先處理,然後將CLIP的預先處理結果存入Tair中,根據Vector提供的近鄰檢索功能,實現高效的圖文多模態檢索。更多關於Tair Vector的資訊,請參見Vector

方案概述

  1. 下載圖片資料。

    本樣本使用的測試資料:

    • 圖片:為開源的寵物圖片資料集,包含7000多張各種形態的寵物圖片。

    • 文本:“狗”、“白色的狗”和“奔跑的白色的狗”。

  2. 串連Tair執行個體,具體實現可參見範例程式碼中的get_tair函數。

  3. 在Tair中分別為圖片(Images)和文字(Texts)建立向量索引,具體實現可參見範例程式碼中的create_index函數。

  4. 寫入圖片和文本資料。

    通過CLIP模型對圖片、文本資料進行預先處理,並通過Vector的TVS.HSET命令,將其名稱和其特徵資訊存入Tair中。具體實現可參見範例程式碼中的insert_images函數(寫入圖片)和upsert_text函數(寫入文本)。多模檢索..jpeg

  5. 進行多模態查詢。

    • 以文搜圖

      通過CLIP模型將待查詢的常值內容進行預先處理,然後通過Vector的TVS.KNNSEARCH命令,查詢Tair資料庫與該文本描述最為相似的圖片。具體實現可參見範例程式碼中的query_images_by_text函數。

    • 以圖搜文

      通過CLIP模型將待查詢的圖片內容進行預先處理,然後通過Vector的TVS.KNNSEARCH命令,查詢Tair資料庫中與該圖片最為吻合的文本。具體實現可參見範例程式碼中的query_texts_by_image函數。

    說明
    • 查詢的文本或圖片無需存入TairVector中。

    • TVS.KNNSEARCH支援指定返回的結果數量(topK),返回的相似距離(distance)越小,表示相似性越高。

樣本代碼

本樣本的Python版本為3.8,且需安裝Tair-py、torch、Image、pylab、plt、CLIP等依賴,其中Tair-py的快捷安裝命令為:pip3 install tair

# -*- coding: utf-8 -*-
# !/usr/bin/env python
from tair import Tair
from tair.tairvector import DistanceMetric
from tair import ResponseError

from typing import List
import torch
from PIL import Image
import pylab
from matplotlib import pyplot as plt
import os
import cn_clip.clip as clip
from cn_clip.clip import available_models


def get_tair() -> Tair:
    """
    該方法用於串連Tair執行個體。
    * host:Tair執行個體串連地址。
    * port:Tair執行個體的連接埠號碼,預設為6379。
    * password:Tair執行個體的密碼(預設帳號)。若通過自訂帳號串連,則密碼格式為“username:password”。
    """
    tair: Tair = Tair(
        host="r-8vbehg90y9rlk9****pd.redis.rds.aliyuncs.com",
        port=6379,
        db=0,
        password="D******3",
        decode_responses=True
    )
    return tair


def create_index():
    """
    建立儲存圖片、文本向量的Vector索引:
    * 圖片Key名稱為"index_images"、文本Key名稱為"index_texts"。
    * 向量維度為1024。
    * 計算向量距離函數為IP。
    * 索引演算法為HNSW。
    """
    ret = tair.tvs_get_index("index_images")
    if ret is None:
        tair.tvs_create_index("index_images", 1024, distance_type="IP",
                              index_type="HNSW")
    ret = tair.tvs_get_index("index_texts")
    if ret is None:
        tair.tvs_create_index("index_texts", 1024, distance_type="IP",
                              index_type="HNSW")


def insert_images(image_dir):
    """
    您需要輸入圖片的路徑,該方法會自動遍曆路徑下的圖片檔案。
    同時,該方法會調用extract_image_features方法(通過CLIP模型對圖片檔案進行預先處理,並返回圖片的特徵資訊),並行將返回的特徵資訊存入TairVector中。
    存入Tair的格式為:
    * 向量索引名稱為“index_images”(固定)。
    * Key為圖片路徑及其檔案名稱,例如“test/images/boxer_18.jpg”。
    * 特徵資訊為1024維向量。
    """
    file_names = [f for f in os.listdir(image_dir) if (f.endswith('.jpg') or f.endswith('.jpeg'))]
    for file_name in file_names:
        image_feature = extract_image_features(image_dir + "/" + file_name)
        tair.tvs_hset("index_images", image_dir + "/" + file_name, image_feature)


def extract_image_features(img_name):
    """
    該方法將通過CLIP模型對圖片檔案進行預先處理,並返回圖片的特徵資訊(1024維向量)。
    """
    image_data = Image.open(img_name).convert("RGB")
    infer_data = preprocess(image_data)
    infer_data = infer_data.unsqueeze(0).to("cuda")
    with torch.no_grad():
        image_features = model.encode_image(infer_data)
    image_features /= image_features.norm(dim=-1, keepdim=True)
    return image_features.cpu().numpy()[0]  # [1, 1024]


def upsert_text(text):
    """
    您需要輸入需儲存的文本,該方法會調用extract_text_features方法(通過CLIP模型對文本進行預先處理,並返迴文本的特徵資訊),並行將返回的特徵資訊存入TairVector中。
    存入Tair的格式為:
    * 向量索引名稱為“index_texts”(固定)。
    * Key為常值內容,例如“奔跑的狗”。
    * 特徵資訊為1024維向量。
    """
    text_features = extract_text_features(text)
    tair.tvs_hset("index_texts", text, text_features)


def extract_text_features(text):
    """
    該方法將通過CLIP模型對文本進行預先處理,並返迴文本的特徵資訊(1024維向量)。
    """
    text_data = clip.tokenize([text]).to("cuda")
    with torch.no_grad():
        text_features = model.encode_text(text_data)
    text_features /= text_features.norm(dim=-1, keepdim=True)
    return text_features.cpu().numpy()[0]  # [1, 1024]


def query_images_by_text(text, topK):
    """
    該方法將用於以文搜圖。
    您需要輸入待搜尋的常值內容(text)和返回的結果數量(topK)。
    該方法會通過CLIP 模型將待查詢的常值內容進行預先處理,然後通過Vector的TVS.KNNSEARCH命令,查詢Tair資料庫中與該文本描述最為相似的圖片。
    將返回靶心圖表片的Key名稱和相似距離(distance),其中相似距離(distance)越小,表示相似性越高。
    """
    text_feature = extract_text_features(text)
    result = tair.tvs_knnsearch("index_images", topK, text_feature)
    for k, s in result:
        print(f'key : {k}, distance : {s}')
        img = Image.open(k.decode('utf-8'))
        plt.imshow(img)
        pylab.show()


def query_texts_by_image(image_path, topK=3):
    """
    該方法將用於以圖搜文。
    您需要輸入待查詢圖片的路徑和返回的結果數量(topK)。
    該方法會通過CLIP 模型將待查詢的圖片內容進行預先處理,然後通過Vector的TVS.KNNSEARCH命令,查詢Tair資料庫中與該圖片最為吻合的文本。
    將返回目標文本的Key名稱和相似距離(distance),其中相似距離(distance)越小,表示相似性越高。
    """
    image_feature = extract_image_features(image_path)
    result = tair.tvs_knnsearch("index_texts", topK, image_feature)
    for k, s in result:
        print(f'text : {k}, distance : {s}')

if __name__ == "__main__":
    # 串連Tair資料庫,並分別建立儲存圖片和文本的向量索引。
    tair = get_tair()
    create_index()
    # 載入Chinese-CLIP模型。
    model, preprocess = clip.load_from_name("RN50", device="cuda", download_root="./")
    model.eval()
    
    # 例如寵物圖片資料集的路徑為“/home/CLIP_Demo”,寫入圖片資料。
    insert_images("/home/CLIP_Demo")
    # 寫入文本樣本資料("狗"、"白色的狗"、"奔跑的白色的狗")。
    upsert_text("狗")
    upsert_text("白色的狗")
    upsert_text("奔跑的白色的狗")

    # 以文搜圖,查詢最符合文本"奔跑的狗"的三張圖。
    query_images_by_text("奔跑的狗", 3)
    # 以圖搜文,指定圖片路徑,查詢比較符合圖片描述的文本。
    query_texts_by_image("/home/CLIP_Demo/boxer_18.jpg",3)

結果展示

  • 以文搜圖,查詢最符合文本"奔跑的狗"的三張圖,結果如下。

    以文搜圖,奔跑的狗..jpeg

  • 以圖搜文,指定查詢圖片如下。

    奔跑的狗(搜尋圖)..jpeg

    查詢結果如下。

    {
      "results":[
        {
          "text":"奔跑的白色的狗",
          "distance": "0.4052203893661499"
        },
        {
          "text":"白色的狗",
          "distance": "0.44666868448257446"
        },
        {
          "text":"狗",
          "distance": "0.4553511142730713"
        }
      ]
    }

總結

Tair作為純記憶體資料庫,內建了HNSW等索引演算法加快了檢索的速度。

採用CLIP與Tair Vector進行多模態檢索,既可以進行以文搜圖(例如商品推薦情境),也可以進行以圖搜文(例如寫作情境)。同時,您也可以採用其他的模型取代CLIP模型,實現以文本搜尋視頻、音頻等更多模態的檢索功能。