利用openai/CLIP模型实现相似图片检索

利用openai/CLIP模型实现相似图片检索

引用

CSDN

1.

https://blog.csdn.net/speedlike/article/details/143806273

CLIP（对比语言-图像预训练）模型是一种通过大量语义相同的（图片，文字）对进行预训练的神经网络，由文本编码器（Text Encoder）和图像编码器（Image Encoder）组成。本文将介绍如何利用CLIP模型实现相似图片检索功能。

CLIP模型原理

CLIP模型的基本原理是让Text Encoder和Image Encoder输出的数字向量在语义空间上靠近。例如，用Text Encoder把文字“可爱的小狗”生成向量emb_t1，把文字“可爱的小猫”生成向量emb_t2，用Image Encoder把一张小狗照片生成向量emb_i1，则在向量空间里面，与emb_i1（即小狗图片的向量表示）距离更近的是emb_t1，而不是emb_t2。

官网的应用DEMO利用clip实现了zero-shot（零训练）的图片分类，在图片集ImageNet上获得了和模型ResNet50差不多的成绩。Github上的官方项目链接：https://github.com/openai/CLIP

实现相似图片检索

利用CLIP的Image Encoder对图片集生成embedding，再利用K临近算法查找向量空间中相似的向量，即可实现相似图片检索（以图找图）。

基本思路

1. 用Image Encoder对图片集中所有照片生成embedding
2. 用hnswlib库将1生成的embedding创建索引index
3. 用Image Encoder对被检索照片生成embedding，记为emb_q
4. 在index中用knn算法找出top_k个与emb_q最相似的embedding（即相似图片的索引编号）

代码实现

input_path = r'D:\mike\test3'#图片库文件夹，任意数量的jpg文件
#生成所有图片文件的embedding
from tqdm import tqdm
import os
from PIL import Image
# 遍历目录下的所有文件（不含子目录），并显示进度
embeddings = []
files = [f for f in os.listdir(input_path) if f.lower().endswith(('.png', '.jpg', '.jpeg','.bmp', '.tiff','.jfif'))]
for filename in tqdm(files, desc='encoding img files'):
    file_path = os.path.join(input_path, filename)
    #过滤小于10k的文件
    #if os.path.getsize(file_path) < 10240:
    #    continue
        
    image = preprocess(Image.open(file_path)).unsqueeze(0).to(device)
    
    with torch.no_grad():
        image_features = model.encode_image(image).squeeze()
    embeddings.append(image_features.cpu().numpy())
    #break #test run, encode only 1 image
assert len(embeddings)==len(files)
#prepare image encoder
import torch
import clip
device = "cuda" if torch.cuda.is_available() else "cpu"
#首次运行会自动下载模型文件，大小300M+
model, preprocess = clip.load("ViT-B/32", device=device)
#创建索引
import hnswlib
index = hnswlib.Index(space='cosine', dim=512)#clip model生成的embedding长度是512
index.init_index(max_elements=len(embeddings), ef_construction=512, M=64)
index.add_items(embeddings, list(range(len(embeddings))))
top_k = 8 # 返回top n个相似的图片
#要查询的图片文件名,默认在图库目录中
#如果要用图库之外的图片来查询，需要单独为target图片生成embedding
query_jpg_filename = 'target.jpg'
#检索出图片文件名对应的embedding index
jpg_idx = files.index(query_jpg_filename)
##查询
query_emb = embeddings[jpg_idx]
simi_jpg_ids,distances = index.knn_query(query_emb, k=top_k)
#print(simi_jpg_ids)
#print(distances)
#显示前8张图片
import matplotlib.pyplot as plt
images_paths=[]
for idx in simi_jpg_ids[0]:
    images_paths.append(os.path.join(input_path, files[idx]))
# 创建一个2行4列的布局
fig, axs = plt.subplots(2, 4, figsize=(12, 8))  # 设置整个figure的大小
# 循环遍历图片路径列表
for i, image_path in enumerate(images_paths):
    # 加载图片
    image = plt.imread(image_path)
    
    # 根据索引确定放置图片的位置
    row, col = divmod(i, 4)
    
    # 在相应的位置显示图片
    axs[row, col].imshow(image)
    axs[row, col].axis('off')  # 关闭坐标轴 
    # 在图片下方添加文字(不支持汉字），写上knn算法返回的向量距离
    axs[row, col].text(0, 0, f'{distances[0][i]:.4f}', 
                       fontsize=16, color='red', ha='center', va='center')
# 调整布局以适应所有图片
plt.tight_layout()
# 显示图像
plt.show()

实际应用效果

实际应用效果见下面图片，每组图片的第1张是被检索的图片，另外7张是找出的相似的图片。

总结

利用CLIP的图片编码器，可实现图片查重的小工具，用于整理图片集，也可应用在图文社区产品中实现基于图片语义的内容推荐。类似的模型还有Open Clip，效果与CLIP相近，最近两年多模态开源模型发展迅猛，利用多模态LLM也可实现图片embedding，并且在语义上更精准，例如可认出图片中的明星是谁，但这些模型动则需要10多G显存，不如CLIP短小精悍。