查询图片（根据自己的图片库）

准备图片数据

可使用图片数据：https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/datasets/unsplash-25k-photos.zip

也可使用自己的图片

图片数据处理

## 获取文件列表
def get_flist(path, ext, ext2=None):
    results = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith(ext):
                results.append(os.path.join(root, f))
            if ext2 is not None and f.endswith(ext2):
                results.append(os.path.join(root, f))
    return results

## 获取图片库信息
def get_img(emb_file, img_path = None):  
    ## 判断是否存在处理好的文件（图片库信息）；有着直接读取，无着编码并存储         
    if os.path.exists(emb_file):
        with open(emb_file, 'rb') as fIn:
            img_names, img_emb = pickle.load(fIn)  
        print("Images:", len(img_names))
        return img_names, img_emb
    else:
        # 图片处理模型
        img_model = SentenceTransformer('clip-ViT-B-32')
    
        img_names = get_flist(img_path, ".jpg", ".png")
        print("Load Images:", len(img_names))
        # 只存储相对路径（便以图片地址改变）
        img_names2 = []
        for img in img_names:
            img_names2.append(img.replace(img_path, ""))
        img_emb = img_model.encode([Image.open(filepath) for filepath in img_names], batch_size=128, convert_to_tensor=True, show_progress_bar=True)
        with open(emb_file, 'wb') as f:
            pickle.dump((img_names2, img_emb), f)
        
        return img_names2, img_emb

查询图片并展示

# 查询
def search(model, img_emb, query, k=3):
    query_emb = model.encode([query], convert_to_tensor=True, show_progress_bar=False)
    hits = util.semantic_search(query_emb, img_emb, top_k=k)[0]
    print("Query:", query)

    return hits

# 展示图片
def image_show(hits, img_names, img_folder):
    for hit in hits:
        img_file = img_folder + img_names[hit['corpus_id']]
        img = cv2.imread(img_file)
        cv2.imshow(str(hit['score']), img)

    if len(hits) > 0:
        cv2.waitKey(0)

# 加载模型
def get_model():
    model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')
    return model

## 图片库地址
img_path = "E:\\image"
## 图片处理后保存的文件（便于下次直接加载）
emb_file = "E:\\image\\embeddings.pkl"
model = get_model()
img_names, img_emb = get_img(emb_file, img_path=img_path)

hits = search(model, img_emb, "两只猫在睡觉")
print(hits)
image_show(hits, img_names, img_path)

结果：

Images: 24996
Query: 两只猫在睡觉
[{'corpus_id': 517, 'score': 0.30299052596092224}, {'corpus_id': 16452, 'score': 0.30096572637557983}, {'corpus_id': 9728, 'score': 0.29650789499282837}]