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

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)              

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)

#加速版的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

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

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()

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()