读取图片测试_基于迁移学习的图片搜索（图片向量搜索）

只开源可执行代码

开篇

一直想写一篇博文，开启我的知乎之旅，今天心血来潮，想起之前做的一个图片搜索的调研，随写一篇简文和大家分享一下，low，low的，希望能对大家有所帮助。

导读

本博文设计到的技术是：图片的迁移特征提取（resnet50），图片的读取（cv2），向量的搜索（faiss）

目录

1，图片的批量读取及数据说明

2，图片的向量特征提取

3，图片向量搜索

4，数据展示

5，总结

正文

1，图片的批量读取

图片的读取我习惯使用cv2,可以将图片向量化读取并且进行 resize，本文的图片存在百度云盘( 密码:z7z0)，是一个小数据集，用于测试本文代码：

import cv2
import numpy as np
from tqdm import tqdm
import os

def load_train(img_size):
    X = np.zeros((len(os.listdir("jpg/")), img_size, img_size, 3), dtype=np.uint8)
    y = []
    k = 0
    for i in tqdm(os.listdir("jpg/")): #tqdm给载入过程增加了进度条
        if "jpg" in i:
            
            X[k] = cv2.resize(cv2.cvtColor(cv2.imread('jpg/%s' % i), cv2.COLOR_BGR2RGB), (img_size, img_size)) #读入图片 + 转成RGB + resize
            y.append(i)
            k +=1
    return X,y
X_train, y_train = load_train(226)              

其中 img_size 用户可以自定义，X_train存储 图片，y_train存储对应 图片名称，方便后续 对应查找图片；这里插一句，cv中原来也有图片特征提取，但是现在收费了。

本文使用数据展示：

2，图片的向量特征提取

将图片向量化：图片向量化就是提取图片特征向量，使用特征向量代表图片；提取图片特征向量的方案有很多，可以自己根据自己的业务领域训练一个分类器，提取送入全连接层的向量作为vec，可以使用迁移学习直接提取图片向量（使用迁移学习提取向量一般向量维度是固定的，可以通过在全连接层之前加入一个dense层降维）。本文使用迁移学习提取向量：

from keras.models import *
from keras.layers import *
from keras.applications import *

input_tensor = Input((226, 226, 3))
inputs = input_tensor

x = Lambda(resnet50.preprocess_input)(inputs) #preprocess_input函数因预训练模型而异
base_model = ResNet50(input_tensor=x, weights='imagenet', include_top=False)
x = base_model(x)
outputs = GlobalAveragePooling2D()(x)
model = Model(inputs, outputs)

train_features = model.predict(X_train) # X_train 是上方的图片读入

train_features.shape

# 提取成功后输出 ：(1362, 2048)

Input 维度可以自定，使用不同维度效果会有一定区别，但总体差异不大；初次使用 resnet50模型会有一个下载过程，可以手动下载，方法自行查阅；

resnet50的网络结构式：

图片网上找的，就为了增加个文章长度 。最终提取的 flatten层 2024维作为图片表示向量。

当然你也可以使用其他的网络，keras自带很多训练好的权重 如：vgg19，Xception等，我也没有测试哪个效果更好，有兴趣可以测试一下给我留言。

3，图片向量搜索

向量搜索原理简单，无非是计算向量的距离，常用的有 欧式距离，余弦距离，相信能看到这个位置的同学肯定不会不知道这些，这里就不 徒增公式了。牛逼的公司就是给我们这些菜鸟提供工具的，facebook开源的faiss就是如此，更开心的是他有python版，里面集成了计算向量相似度的多种方法，有需要的可以去看官方文档 python安装命令是 pip3 install faiss-cpu （也可以有gpu版）。这里的向量搜索就是使用faiss：

#加速版的cosine_similarity的计算
import time, faiss
from faiss import normalize_L2

d = 2048                           # dimension
print(train_features.shape)
nb = train_features.shape[0]                    # database size         # make reproducible
normalize_L2(train_features)

nlist = 100 # 聚类中心的个数
k = 5 #邻居个数  就是你想输出top几个
quantizer = faiss.IndexFlatIP(d)  # the other index，需要以其他index作为基础

index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_INNER_PRODUCT)
# by default it performs inner-product search
assert not index.is_trained
index.train(train_features)
assert index.is_trained
index.nprobe = 300  # default nprobe is 1, try a few more
index.add(train_features)  # add may be a bit slower as well
t1=time.time()
def IndexIVFFlat(training_vectors):
    D, I = index.search(training_vectors, k)  # actual search
    t2 = time.time()
    print('faiss kmeans result times {}'.format(t2-t1))
    

    # 用来输出对应的index 和 相似度；index用来寻找 y_train中对应的 图片名称
    for i,d in zip(I[0],D[0]):
        print(y_train[i],d)
    return I,D

上边的代码是 将 得到的图片向量矩阵 送入 faiss 训练，方便搜索，效率很高，就是占内存比较大（当然是说你图片上kw时，因为faiss需要被搜索的向量全部住内存，这也为什么我说，你可以在迁移学习时增加一个dense层来降低 图片向量维度的原因）

测试一下：

i = 100 # 随便拿一张做测试
IndexIVFFlat(train_features[i:i+1])

faiss kmeans result times 132.28123712539673
image_0838.jpg 0.9999996
image_0866.jpg 0.9999946
image_0841.jpg 0.9163162
image_0802.jpg 0.9163011
image_0862.jpg 0.9110917

4，数据展示

得到 相似结果可以简单展示一下；

from matplotlib import pyplot as plt

p= y_train[i] # i是上边的测试输入index，y_train[i]是对应 的图片名称

img = cv2.imread('jpg/' + p, 0)
plt.imshow(img, cmap='gray')
plt.savefig('plt.png')
plt.show()

展示第三个最相似的，因为第一个是它本身，毕竟我的测试数据是在 x_train中取得吗。

from matplotlib import pyplot as plt

p= "image_0841.jpg"

img = cv2.imread('jpg/' + p, 0)
plt.imshow(img, cmap='gray')
plt.savefig('plt.png')
plt.show()

嗯，至少我觉还挺相似的哈哈。。。。。

5，总结

本人是做nlp的，图片是个爱好，有问题欢迎大家留言指正，交流和合作可以留言。关于faiss有个提醒，就是 使用新图片 通过特征提取得到向量后与 已faiss训练的图片计算余弦相似度时你可能会发现相似度大于1，这是因为，你的向量没有 做归一化，就是说 除以 向量本身的模长，不懂得可留言！欢迎加微信交流合作