AI驱动的相似图像搜索

你知道当你半夜睡不着觉时,那些奇怪的想法是如何随机出现在你脑海里的吗?我躺在床上,突然,我突然想到——“等等,印度尼西亚和摩纳哥的国旗一模一样吗?”好吧,这就是几周前的我。前一刻,我还在浏览随机的维基百科文章,下一刻,我正全身心投入一个使用人工智能比较国旗的项目。

那么,我到底做了什么?我决定使用不同的人工智能模型,如 ViT、CLIP、BLIP、EfficientNet、DINO-v2 和值得信赖的旧 VGG16 来比较图像并查看它们的相似性。它从国旗开始,但让我陷入了使用深度学习模型探索图像相似性这一更广泛概念的兔子洞。

1、直觉:图像作为空间中的向量

图像相似性搜索的核心是一个简单的想法:图像可以表示为高维空间中的向量。当两幅图像相似时,它们的向量应该在这个空间中占据相似的位置。我们可以通过测量角度(或余弦相似度)来确定这些向量的相似程度。如果角度较小,则图像接近(相似)。如果角度较大(不同),则图像相距较远。这就像在公寓大楼里寻找邻居,但方式要抽象得多。

图像的向量嵌入

2、了解AI模型

我使用了几种深度学习模型,每种模型都有其优点和特点:

  • CLIP(对比语言-图像预训练):由 OpenAI 构建,它学习将图像与文本进行匹配。对于我们的相似性搜索来说,这是一个不错的选择。
  • ViT(视觉变换器):ViT 通过将图像视为序列来彻底改变图像处理,类似于变换器处理文本的方式。
  • BLIP:一种视觉语言模型,专注于对齐视觉和文本内容。
  • EfficientNet:该模型以其效率而闻名,非常适合图像识别任务。
  • DINO:一种自监督变换器模型,擅长从图像中学习特征。
  • VGG16:一种经典的卷积神经网络 (CNN),已经存在多年,并且在图像识别任务中仍然占有一席之地。

我不会在这里深入探讨每个模型的架构细节,但如果你有兴趣了解更多有关它们的信息,请查看此处的精彩文章。

3、数据准备

我从维基百科上抓取了国旗的图片,将世界各国国旗变成了数据集。

import pandas as pd
flags_df = pd.read_csv('national_flags.csv')  
print(flags_df)
从 Wikipedia 抓取的数据

4、特征提取

一旦我得到了标志,真正的乐趣就开始了:提取特征。每个模型都获取标志图像并将其转换为特征向量。这就像将图像转换为封装其特征的数字列表。

我将使用 Huggingface Transformer 库来提取此实验中的嵌入。

4.1 EfficientNet

通过平均最后一个隐藏层输出的空间维度来提取标志特征,重点关注细粒度模式。

image_processor = AutoImageProcessor.from_pretrained("google/efficientnet-b7")
model = EfficientNetModel.from_pretrained("google/efficientnet-b7")

# prepare input image
inputs = image_processor(img, return_tensors='pt')

with torch.no_grad():
    outputs = model(**inputs, output_hidden_states=True)

embedding = outputs.hidden_states[-1]
embedding = torch.mean(embedding, dim=[2,3])

4.2 ViT

使用其转换器架构中第一个标记的最后一个隐藏状态,捕获局部和全局视觉特征。

image_processor = AutoImageProcessor.from_pretrained("google/vit-large-patch16-224-in21k")
model = ViTModel.from_pretrained("google/vit-large-patch16-224-in21k")

# prepare input image
inputs = image_processor(img, return_tensors='pt')

with torch.no_grad():
    outputs = model(**inputs)
embedding = outputs.last_hidden_state
embedding = embedding[:, 0, :].squeeze(1)

4.3 DINO-v2

通过专注于自监督学习来生成嵌入,利用第一个标记来捕获以对象为中心的细节。

image_processor = AutoImageProcessor.from_pretrained('facebook/dinov2-base')
model = AutoModel.from_pretrained('facebook/dinov2-base')

# prepare input image
inputs = image_processor(img, return_tensors='pt')

with torch.no_grad():
    outputs = model(**inputs)
embedding = outputs.last_hidden_state
embedding = embedding[:, 0, :].squeeze(1)

4.4 CLIP

结合图像和文本嵌入,使用图像特征来理解视觉概念以及来自配对文本的上下文数据。

image_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

# prepare input image
inputs = image_processor(images=img, return_tensors='pt', padding=True)

with torch.no_grad():
    embedding = model.get_image_features(**inputs) 

4.5 BLIP-2

采用视觉语言模型,通过查询聚焦转换器(Q-Former)提取特征来捕捉图像语义和关系。

image_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
model = Blip2Model.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16)

inputs = image_processor(images=img, return_tensors='pt', padding=True)
print('input shape: ', inputs['pixel_values'].shape)

with torch.no_grad():
    outputs = model.get_qformer_features(**inputs)
embedding = outputs.last_hidden_state
embedding = embedding[:, 0, :].squeeze(1)

4.6 VGG16

一种 CNN 模型,通过应用多层卷积层来输出标志嵌入,强调分层图像表示。

model = models.vgg16(pretrained=True) 
model.eval()  # Set the model to evaluation mode

batch_t = torch.unsqueeze(img, 0)

with torch.no_grad():
    embedding = model(batch_t)

4.7 提取所有旗帜

接下来,我们将每个模型应用于 DataFrame 中的旗帜图像,提取它们的特征嵌入,作为相似性分析的基础。

# Extract features for all flags
flags_df['features'] = flags_df['Flag Image'].apply(extract_features)

5、使用 FAISS 进行余弦相似度计算

现在我们已经锁定并加载了旗帜嵌入,我们将利用余弦相似度来衡量旗帜嵌入的相似度。余弦相似度比较两个向量的方向,使其能够有效地根据模式而不是幅度识别关系。这种方法在分析旗帜嵌入时特别有用,因为重点是特征向量的相对形状和设计元素,而不是绝对大小。

要实现这一点:

  • 规范化:每个特征向量都标准化为单位长度,因此可以将余弦相似度计算为向量的点积。这确保了相似度反映了向量之间的角度。
  • FAISS 用于相似度搜索:我们利用 FAISS(一个针对高效相似度搜索进行了优化的库)根据其标准化的旗帜嵌入找到前 K 个最相似的国家。这允许在大型旗帜图像数据集中进行快速且可扩展的比较。

该方法提供了更具上下文感知的相似性搜索,捕捉旗帜图案的复杂细节,并确保顶部相似的旗帜具有视觉意义。

def clean_feature_string(feature_str):
    cleaned_str = re.sub(r'[\[\]]', '', feature_str)  # Remove brackets
    cleaned_values = np.fromstring(cleaned_str, sep=' ')  # Parse values into numpy array
    return cleaned_values

# Function to get top K similar countries using FAISS
def get_top_k_similar_countries(input_country, df, k=5):
    countries = df['Country'].values
    features = np.array([clean_feature_string(f) for f in df['features'].values])
    
    # Find the index of the input country
    try:
        input_idx = list(countries).index(input_country)
    except ValueError:
        return f"Country '{input_country}' not found in the dataset."
    
    input_embedding = features[input_idx].reshape(1, -1)

    # Normalize the feature vectors for cosine similarity
    features_normalized = features / np.linalg.norm(features, axis=1, keepdims=True)

    # Create a FAISS index for similarity search
    dim = features.shape[1]
    index = faiss.IndexFlatIP(dim)  
    
    # Add all features to the FAISS index
    index.add(features_normalized)
    
    # Search for the top K most similar countries
    distances, top_k_idx = index.search(input_embedding, k+1)  # k+1 to exclude the country itself
    
    # Return top K countries with their similarity scores
    return [(countries[i], distances[0][j]) for j, i in enumerate(top_k_idx[0]) if i != input_idx]

# Display top 5 similar flags 
top_5_countries = get_top_k_similar_countries(country, k=5)

for idx, (country, score) in enumerate(top_5_countries):
    # Load the flag image for each country from the local folder
    img = load_local_image(country)
    display(img)

6、评估:国旗相似性和偏见

事情从这里开始变得有趣——也有点个人化。图像相似性是一个非常主观的领域。

在我看来,两面相似的国旗可能不会引起别人的共鸣。在评估过程中,我注意到自己的偏见悄悄潜入其中,影响了我对结果的看法。

为了解决这个问题,我在《大英百科全书》上找到了一篇比较不同国家相似国旗的文章,并决定将其作为基准。挑战是什么?我能用人工智能复制这些结果吗?

为了测试它,我在所有模型中运行了同一组国旗图像——并比较了它们的最佳推荐。过程如下:

测试 1:乍得和罗马尼亚

所有模型都返回“罗马尼亚”作为最佳匹配。

与 乍得 相似的旗帜
测试 2:塞内加尔和马里

只有 CLIP 和 VGG16 正确识别了“马里”,但令人惊讶的是,所有模型都选择了“喀麦隆”作为最接近的匹配项(这让我们比《大英百科全书》领先了一步)。

与塞内加尔相似的国旗

测试 3:印度尼西亚和摩纳哥

啊,这对组合让我彻夜难眠。所有模型都返回了“摩纳哥”,尽管并非每次都是最高得分。

与印度尼西亚相似的国旗
测试 4:澳大利亚和新西兰

所有模型都正确识别了新西兰,但引起我注意的是,一些模型还将图瓦卢标记为匹配项——这是一个有趣的发现!

与澳大利亚相似的旗帜

作为家庭作业(为了避免这成为一场滚动马拉松),我将留下一些配对供您自己测试:

  • 爱尔兰和科特迪瓦(象牙海岸)
  • 挪威和冰岛
  • 委内瑞拉、厄瓜多尔和哥伦比亚
  • 卢森堡和荷兰
  • 斯洛文尼亚、俄罗斯和斯洛伐克

我已经在 Streamlit 上部署了该应用程序,所以看看你的结果是否匹配。

7、注意:相似性模型中的语境偏差

某些模型(如 CLIP)可能会无意中根据语境关系而非视觉相似性对国家进行分组。例如,印度和巴基斯坦在相似性搜索中可能被放在一起,不是因为它们的国旗设计,而是因为它们的地缘政治联系以及在全球新闻和讨论中经常一起提及。同样,模型可能会将以色列和巴勒斯坦或乌克兰和俄罗斯归为一组,即使它们的国旗有很大不同,因为这些国家在政治话语、媒体报道和历史背景方面往往有联系。

相似性模型中的语境偏差

之所以发生这种情况,是因为在图像和文本上训练的模型(如 CLIP)可以发现这些更广泛的关联,而不仅仅是视觉数据。这一现象凸显了模型不仅可能反映视觉相似性,还可能反映训练数据中存在的语境偏差。

8、结束语

这个项目是对不同 AI 模型如何处理图像相似性的有趣探索。每个模型都有独特的优势,虽然没有一个模型是完美的,但它们结合在一起提供了关于机器如何解释视觉数据的有趣见解。无论你处理的是旗帜、徽标还是任何其他类型的图像,这些模型都可以适应无数应用程序中的特征提取、图像比较和相似性搜索等任务。如果您对细节感兴趣,请查看GitHub 存储库中的代码和结果。

最终,AI驱动的相似性搜索除了旗帜之外还有广泛的潜在应用——想想徽标、艺术品、设计图案,甚至更复杂的任务,如面部识别或医学成像。可能性无穷无尽。


原文链接:Build an AI Image Similarity Search with Transformers — ViT, CLIP, DINO-v2, and BLIP-2

BimAnt翻译整理,转载请标明出处