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

# -*- 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)