如何实现以图搜图

爱被打了一巴掌 2024-03-17 21:10 52阅读 0赞

## 一、前言 ##

在许多搜索引擎中，都内置了以图搜图的功能。以图搜图功能，可以极大简化搜索工作。今天要做的就是实现一个以图搜图引擎。

我们先来讨论一下以图搜图的难点，首当其冲的就是如何对比图片的相似度？怎么样的图片才叫相似？人可以一眼判断，但是计算机却不一样。图片以数字矩阵的形式存在，而相似度的比较也是比较矩阵的相似度。但是这有一些问题。

第二个问题就是大小问题，图片的大小通常是不一样的，而不同大小的矩阵也无法比较相似度。不过这个很好解决，直接修改图片尺寸即可。

第三个问题则是像素包含的信息非常有限，无法表达抽象信息。比如画风、物体、色调等。

根据上面描述，我们现在要解决两个问题：用什么信息替换像素信息、怎么计算相似度。下面一一解决。在开始前，我们先实现一个简易的以图搜图功能。

## 二、简易以图搜图实现 ##

### 2.1 如何计算相似度 ###

首先来讨论一下直接使用像素作为图像的表示，此时我们应该如何完成以图搜图的操作。一个非常简单的想法就是直接计算两个图片的几何距离，假如我们目标图片为target，图库中的图片为source，几何距离的计算如下：

distance = sum\[(target−source)2\]distance = \\sqrt\{sum\[(target - source)^2\]\}distance = sum\[(target−source)2\]

然后把距离最小的n个图片作为搜索结果。

这个方法看起来不可靠，但是实际使用时也会有不错的结果。如果图库图片本身不是非常复杂，比如动漫头像，那么这种方式非常简单有效，而其它情况下结果会比较不稳定。

### 2.2 基于几何距离的图片搜索 ###

基于几何距离的图片搜索实现步骤如下：

1.  把图片修改到同一尺寸，如果尺寸不同则无法计算几何距离
2.  选定一个图片作为目标图片，即待搜索图片
3.  遍历图库，计算几何距离，并记录到列表
4.  对列表排序，获取几何距离最小的n张图片

这里使用蜡笔小新的图片作为图库进行搜索，下面是图片的一些示例：

![format_png][]

部分图片有类似的风格，我们希望能根据一张图片找到类似风格的图片。实现代码如下：

import os
    import cv2
    import random
    
    import numpy as np
    
    base_path = r"G:\datasets\lbxx"
    # 获取所有图片路径
    files = [os.path.join(base_path, file) for file in os.listdir(base_path)]
    # 选取一张图片作为目标图片
    target_path = random.choice(files)
    target = cv2.imread(target_path)
    h, w, _ = target.shape
    distances = []
    # 遍历图库
    for file in files:
        # 读取图片，转换成与目标图片同一尺寸
        source = cv2.imread(file)
        if not isinstance(source, np.ndarray):
            continue
        source = cv2.resize(source, (w, h))
        # 计算几何距离，并加入列表，这里没有开方
        distance = ((target - source) ** 2).sum()
        distances.append((file, distance))
    # 找到相似度前5的图片，这里拿了6个，第一个是原图
    distances = sorted(distances, key=lambda x: x[-1])[:6]
    imgs = list(map(lambda x: cv2.imread(x[0]), distances))
    result = np.hstack(imgs)
    cv2.imwrite("result.jpg", result)

下面是一些比较好搜索结果，其中最左边是target，其余10张为搜索结果。

![format_png 1][]

![format_png 2][]

![format_png 3][]

如果换成猫狗图片，下面是一些搜索结果：

![format_png 4][]

![format_png 5][]

### 2.3 存在的问题 ###

上面的实现存在两个问题，其一是像素并不能表示图像的深层含义。搜索结果中经常会返回颜色相似的图片。第二个则是计算复杂度的问题，如果图片大小未224×224，那么图片有150528个像素，计算几何距离会比较耗时。而且在搜索时，需要遍历整个图库，当图库数量较大时，计算量将不可忍受。因此需要对上面的方法进行一些改进。

## 三、改进一，用特征代替像素 ##

### 3.1 图像特征 ###

在表示图片时，就是从基本的像素到手工特征再到深度学习特征。相比之下，用卷积神经网络提取的图像特征有几个有点，具体如下：

1.  具有很强的泛化能力，提取的特征受角度、位置、亮度等的影响会像素和手工特征。
2.  较少的维度，使用ResNet50提取224×224图片的特征时，会返回一个7×7×2048的张量，这比像素数量要少许多。
3.  具有抽象性，相比前面两种，卷积神经网络提取的特征具有抽象性。比如关于图片中类别的信息，这是前面两种无法达到的效果。

在本文我们会使用ResNet50来提取图片特征。

### 3.2 Embedding的妙用 ###

使用ResNet50提取的特征也可以被称为Embedding，也可以简单理解为图向量。Embedding近几年在人工智能领域发挥了巨大潜力，尤其在自然语言处理领域。

#### 3.2.1 关系可视化 ####

早期Embedding主要用于词向量，通过word2vec把单词转换成向量，然后就可以完成一些奇妙的操作。比如单词之间关系的可视化，比如下面这张图：

![format_png 6][]

在图片中可视化了：mother、father、car、auto、cat、tiger六个单词，从图可以明显看出mother、father比较近；car、auto比较近；cat、tiger比较近，这些都与我们常识相符。

#### 3.2.2 关系运算 ####

我们希望训练良好的Embedding每一个维度都有一个具体的含义，比如第一维表示词性，第二维表示情感，其余各个维度都有具体含义。如果能达到这个效果，或者达到近似效果，那么就可以使用向量的计算来计算单词之间的关系。

比如“妈妈-女性+男性≈爸爸”，或者“国王-男性+女性≈皇后”。比如以往要搜索“物理学界的贝多芬是谁”可能得到非常奇怪的结果，但是如果把这个问题转换成“贝多芬-音乐界+物理学界≈?”，这样问题就简单多了。

#### 3.2.3 聚类 ####

当我们可以用Embedding表示图片和文字时，就可以使用聚类算法完成图片或文字的自动分组。在许多手机的相册中，就有自动图片归类的功能。

聚类还可以加速搜索的操作，这点会在后面详细说。

### 3.3 以图搜图改进 ###

下面使用图像特征来代替像素改进以图搜图，代码如下：

import os
    import cv2
    import random
    import numpy as np
    from keras.api.keras.applications.resnet50 import ResNet50
    from keras.api.keras.applications.resnet50 import preprocess_input
    
    w, h = 224, 224
    # 加载模型
    encoder = ResNet50(include_top=False)
    
    
    base_path = r"G:\datasets\lbxx"
    files = [os.path.join(base_path, file) for file in os.listdir(base_path)]
    target_path = random.choice(files)
    target = cv2.resize(cv2.imread(target_path), (w, h))
    
    
    # 提取图片特征
    target = encoder(preprocess_input(target[None]))
    distances = []
    for file in files:
        source = cv2.imread(file)
        if not isinstance(source, np.ndarray):
            continue
            
        # 读取图片，提取图片特征
        source = cv2.resize(source, (w, h))
        source = encoder(preprocess_input(source[None]))
        
    
        distance = np.sum((target - source) ** 2)
        distances.append((file, distance))
    # 找到相似度前5的图片，这里拿了6个，第一个是原图
    distances = sorted(distances, key=lambda x: x[-1])[:6]
    imgs = list(map(lambda x: cv2.imread(x[0]), distances))
    result = np.hstack(imgs)
    cv2.imwrite("result.jpg", result)

这里使用在imagenet上预训练的ResNet50作为特征提取网络，提取的关键操作如下：

1.  加载模型

# 加载ResNet50的卷积层，舍弃全连接部分
    encoder = ResNet50(include_top=False)

1.  图片预处理

# 把图片转换成224×224，并使用ResNet50内置的预处理方法处理
    target = cv2.resize(cv2.imread(target_path), (w, h))
    target = preprocess_input(target[None])

1.  提取特征

# 使用ResNet40网络提取特征
    target = encoder(preprocess_input(target)

下面是改进后的搜索结果：

![format_png 7][]

## 四、改进二，使用聚类改进搜索速度 ##

### 4.1 实现原理 ###

在前面的例子中，我们都是使用线性搜索的方式，此时需要遍历所有图片。搜索复杂度为O(n)，通常可以用树结构来存储待搜索的内容，从而把复杂度降低到O(logn)。这里我们使用更简单的方法，即聚类。

首先我们要做的就是对图片的特征进行聚类，聚成c个簇，每个簇都会对应一个簇中心。簇中心可以认为是一个簇中的平均结构，同一簇中的样本相似度会比较高。

在完成聚类后，我们可以拿到target图片的向量，在c个簇中心中查找target与哪个簇最接近。然后再到当前簇中线性查找最相似的几个图片。

![format_png 8][]

### 4.2 代码实现 ###

代码实现分为下面几个步骤：

1.  把图片转换成向量

这部分代码和前面基本一样，不过这次为了速度快，我们把图像特征存储到embeddings.pkl文件：

import os
    import cv2
    import pickle
    import numpy as np
    import tensorflow as tf
    from keras.api.keras.applications.resnet50 import ResNet50
    from keras.api.keras.applications.resnet50 import preprocess_input
    
    w, h = 224, 224
    # 加载模型
    encoder = ResNet50(include_top=False)
    base_path = r"G:\datasets\lbxx"
    # 获取所有图片路径
    files = [os.path.join(base_path, file) for file in os.listdir(base_path)]
    # 将图片转换成向量
    embeddings = []
    for file in files:
        # 读取图片，转换成与目标图片同一尺寸
        source = cv2.imread(file)
        if not isinstance(source, np.ndarray):
            continue
        source = cv2.resize(source, (w, h))
        embedding = encoder(preprocess_input(source[None]))
        embeddings.append({
            "filepath": file,
            "embedding": tf.reshape(embedding, (-1,))
        })
    with open('embeddings.pkl', 'wb') as f:
        pickle.dump(embeddings, f)

1.  对所有向量进行聚类操作

这里可以使用sklearn完成：

from sklearn.cluster import KMeans
    with open('embeddings.pkl', 'rb') as f:
        embeddings = pickle.load(f)
    X = [item['embedding'] for item in embeddings]
    kmeans = KMeans(n_clusters=500)
    kmeans.fit(X)
    preds = kmeans.predict(X)
    for item, pred in zip(embeddings, preds):
        item['cluster'] = pred
    joblib.dump(kmeans, 'kmeans.pkl')
    with open('embeddings.pkl', 'wb') as f:
        pickle.dump(embeddings, f)

如果图片数量比较多的话，这部分操作会比较耗时。然后调用kmeans.predict方法就可以知道某个图片属于哪个簇，这个也可以事先存储。

1.  找到输入图片最近的簇中心

在训练完成后，就可以拿到所有簇中心：

kmeans.cluster_centers_

现在要做的就是找到与输入图片最近的簇中心，这个和前面的搜索一样：

# 查找最近的簇
    closet_cluster = 0
    closet_distance = sys.float_info.max
    for idx, center in enumerate(centers):
        distance = np.sum((target.numpy() - center) ** 2)
        if distance < closet_distance:
            closet_distance = distance
            closet_cluster = idx

1.  在当前簇中查找图片

这个和前面也是基本一样的：

distances = []
    for item in embeddings:
        if not item['cluster'] == closet_cluster:
            continue
        embedding = item['embedding']
        distance = np.sum((target - embedding) ** 2)
        distances.append((item['filepath'], distance))
    # 对距离进行排序
    distances = sorted(distances, key=lambda x: x[-1])[:11]
    imgs = list(map(lambda x: cv2.imread(x[0]), distances))
    result = np.hstack(imgs)
    cv2.imwrite("result.jpg", result)

下面是一些搜索结果：

![format_png 9][]

![format_png 10][]

![format_png 11][]

效果还是不错的，而且这次搜索速度快了许多。不过在编码上这种方式比较繁琐，为了让代码更简洁，下面引入向量数据库。

## 五、向量数据库 ##

### 5.1 向量数据库 ###

向量数据库和传统数据库不太一样，可以在数据库中存储向量字段，然后完成向量相似度检索。使用向量数据库可以很方便实现上面的检索功能，而且性能方面会比前面更佳。

向量数据库与传统数据库有很多相似的地方，在关系型数据库中，数据库分为连接、数据库、表、对象。在向量数据库中分别对应连接、数据库、集合、数据。集合中，可以添加embedding类型的字段，该字段可以用于向量检索。

### 5.2 Milvus向量数据库的使用 ###

下面简单说一下Milvus向量数据库的使用，首先需要安装Milvus，执行下面两条执行即可：

wget https://github.com/milvus-io/milvus/releases/download/v2.2.11/milvus-standalone-docker-compose.yml -O docker-compose.yml
    sudo docker-compose up -d

下载完成后，需要连接数据库，代码如下：

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
    connections.connect(host='127.0.0.1', port='19530')

然后创建集合：

def create_milvus_collection(collection_name, dim):
        if utility.has_collection(collection_name):
            utility.drop_collection(collection_name)
    
        fields = [
            FieldSchema(name='id', dtype=DataType.INT64, descrition='ids', max_length=500, is_primary=True,
                        auto_id=True),
            FieldSchema(name='filepath', dtype=DataType.VARCHAR, description='filepath', max_length=512),
            FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, descrition='embedding vectors', dim=dim),
        ]
        schema = CollectionSchema(fields=fields, description='reverse image search')
        collection = Collection(name=collection_name, schema=schema)
    
        # create IVF_FLAT index for collection.
        index_params = {
            'metric_type': 'L2',
            'index_type': "IVF_FLAT",
            'params': {"nlist": 2048}
        }
        collection.create_index(field_name="embedding", index_params=index_params)
        return collection
    
    
    collection = create_milvus_collection('images', 2048)

其中create\_milvus\_collection的第二个参数是embedding的维度，这里传入图片特征的维度。然后把图片特征存储到向量数据库中，这里需要注意维度不能超过32768，但是ResNet50返回的维度超过了这个限制，为此可以用PCA降维或者采用其它方法获取图片embedding。

import pickle
    from sklearn.decomposition import PCA
    
    with open('embeddings.pkl', 'rb') as f:
        embeddings = pickle.load(f)
    X = [item['embedding'] for item in embeddings]
    pca = PCA(n_components=2048)
    X = pca.fit_transform(X)
    for item, vec in zip(embeddings, X):
        item['embedding'] = vec
    with open('embeddings.pkl', 'wb') as f:
        pickle.dump(embeddings, f)
    with open('pca.pkl', 'wb') as f:
        pickle.dump(pca, f)

这样就可以插入数据了，代码如下：

index_params = {
        "metric_type": "L2",
        "index_type": "IVF_FLAT",
        "params": {"nlist": 1024}
    }
    with open('embeddings.pkl', 'rb') as f:
        embeddings = pickle.load(f)
    base_path = r"G:\datasets\lbxx"
    # 获取所有图片路径
    files = [os.path.join(base_path, file) for file in os.listdir(base_path)]
    for item in embeddings:
        collection.insert([
            [item['filepath']],
            [item['embedding']]
        ])

现在如果想要搜索图片，只需要下面几行代码即可：

import os
    import cv2
    import joblib
    import random
    
    import numpy as np
    import tensorflow as tf
    from PIL import Image
    from keras.api.keras.applications.resnet50 import ResNet50
    from keras.api.keras.applications.resnet50 import preprocess_input
    from pymilvus import connections, Collection
    
    pca = joblib.load('pca.pkl')
    w, h = 224, 224
    encoder = ResNet50(include_top=False)
    base_path = r"G:\datasets\lbxx"
    files = [os.path.join(base_path, file) for file in os.listdir(base_path)]
    target_path = random.choice(files)
    target = cv2.resize(cv2.imread(target_path), (w, h))
    target = encoder(preprocess_input(target[None]))
    target = tf.reshape(target, (1, -1))
    target = pca.transform(target)
    
    # 连接数据库，加载images集合
    connections.connect(host='127.0.0.1', port='19530')
    collection = Collection(name='images')
    search_params = {"metric_type": "L2", "params": {"nprobe": 10}, "offset": 5}
    collection.load()
    # 在数据库中搜索
    results = collection.search(
        data=[target[0]],
        anns_field='embedding',
        param=search_params,
        output_fields=['filepath'],
        limit=10,
        consistency_level="Strong"
    )
    collection.release()
    images = []
    for result in results[0]:
        entity = result.entity
        filepath = entity.get('filepath')
        image = cv2.resize(cv2.imread(filepath), (w, h))
        images.append(np.array(image))
    result = np.hstack(images)
    cv2.imwrite("result.jpg", result)

下面是一些搜索结果，整体来看还是非常不错的，不过由于降维的关系，搜索效果可能或略差于前面，但是整体效率要高许多。

![format_png 12][]

![format_png 13][]

![format_png 14][]

## 六、总结 ##

本文我们分享了以图搜图的功能。主要思想就是将图片转换成向量表示，然后利用相似度计算，在图库中查找与之最接近的图片。最开始使用线性搜索的方式，此时查找效率最低。而后使用聚类进行改进，把把图片分成多个簇，把查找分为查找簇和查找最近图片两个步骤，可以大大提高查找效率。

改进后代码变得比较繁琐，于是引入向量数据库，使用向量数据库完成检索功能。这样就完成了整个程序的编写。

[format_png]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/507815e693294c46a6263081709cc75e.png
[format_png 1]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/16c6e71c3d05409e9fd0d6fb7555a1e6.png
[format_png 2]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/ca9bba55d63f4018a50c5c9eafa5e02b.png
[format_png 3]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/26ef0846849e4969bbbb2f8b05c75ed6.png
[format_png 4]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/48cb1a63baa744d38278efad1fef5ae1.png
[format_png 5]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/8c156c3ab5b5469cb36f55def8bc17ec.png
[format_png 6]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/94a61ceabc7748ac80dac2ffccabce5b.png
[format_png 7]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/0ebf8005f23f4b1ba7f712b3eb6e5571.png
[format_png 8]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/932fda4f3cb64d88838a42f1b5a50a55.png
[format_png 9]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/2469dfa821ca4243a527ed23d2dd0141.png
[format_png 10]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/449e689831284a4bafb03a9fa784affa.png
[format_png 11]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/ba9415e3a9644927b8b6e1ee507f780a.png
[format_png 12]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/29f04bd6fd64486bb93e036eb689da8e.png
[format_png 13]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/7f197038b53e49d98c7d92c6cc7fc781.png
[format_png 14]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/04ae608115c542aab4b468424aa878b9.png