利用 inceptionV3 特征简单实现以图搜图

从百度图片搜集一些素材:

分别有：喵x4，汪x4，猪x4，鸡x4

实验的目的

是希望通过输入一张图片，返回与其最相似的图片，达到以图搜图的效果。

代码

首先，导入必要的包

import os
import tarfile
import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
import tensorflow.contrib.slim as slim
import tensorflow.contrib.slim.nets as nets
from six.moves.urllib.request import urlretrieve

# 获取图片
images = []
img_dir = '\Download\img'
for idx, img_name in enumerate(os.listdir(img_dir)):
    img_path = os.path.join(img_dir,img_name)
    print('fetching pic:',img_path)
    img = PIL.Image.open(img_path)
    img = img.resize((299, 299))
    images.append(img)
NUM = len(images)

从网上下载别人训练好的模型，也可以手动下载解压到文件夹。

# 下载 inceptionV3 模型
data_dir = '.'
checkpoint_filename = os.path.join(data_dir, 'inception_v3.ckpt')
if not os.path.exists(checkpoint_filename):
    print('downloading inceptionV3 model...')
    inception_tarball, _ = urlretrieve(
        'http://download.tensorflow.org/models/inception_v3_2016_08_28.tar.gz')
    tarfile.open(inception_tarball, 'r:gz').extractall(data_dir)

# 计算流 （注意：299,299,3,1001是inceptionV3的标准参数
image = tf.Variable(tf.zeros((299, 299, 3)))
preprocessed = tf.multiply(tf.subtract(tf.expand_dims(image, 0), 0.5), 2.0)
arg_scope = nets.inception.inception_v3_arg_scope(weight_decay=0.0)
with slim.arg_scope(arg_scope):
    logits, end_points = nets.inception.inception_v3(
        preprocessed, 1001, is_training=False, reuse=False)
    # logits = logits[:, 1:]  # ignore background class
    # probs = tf.nn.softmax(logits)  # probabilities

#利用 saver 载入网络参数
restore_vars = [
    var for var in tf.global_variables() if var.name.startswith('InceptionV3/')
]
saver = tf.train.Saver(restore_vars)

在 session 中得到所有图片的 inception 特征

with tf.Session() as sess:
    saver.restore(sess, os.path.join(data_dir, 'inception_v3.ckpt')) #加载网络参数
    layer = 'PreLogits'
    features = []
    for img in images:
        img = np.asarray(img).astype(np.float32)/ 255.0
        feature_values = sess.run( end_points, feed_dict={image: img})
        feature = feature_values[layer].squeeze()
        features.append(feature)
    feature_vectors = np.stack(features)

# 计算欧氏距离
# 在这里不做归一化，因为归一化之后和余弦距离等价
distance_euclidean = np.sum(
    np.power(feature_vectors, 2), axis=1, keepdims=True) + np.sum(
        np.power(feature_vectors, 2), axis=1,
        keepdims=True).T - 2 * np.dot(feature_vectors, feature_vectors.T)

# 计算余弦距离
features_norm = feature_vectors / np.linalg.norm(
    feature_vectors, axis=1)[:, np.newaxis]
distance_cosin = np.dot(features_norm, features_norm.T)

到这里，所有图片之间的相似度就已经得到了。

接下来，我们从素材中挑4张图片来展示效果: 找出和输入图片最相似的前6张图片

test_images = [(i,img) for i, img in enumerate(images) if i%4==2]
TEST_NUM = len(test_images)
TOP_NUM = 6

# 测试欧式距离
plt.figure()
for i,(idx_img,plt_img) in enumerate(test_images):
    order_euclidean = np.argsort(distance_euclidean[idx_img])

    plt.subplot(TEST_NUM, TOP_NUM+1, i * (TOP_NUM+1) + 1)
    plt.axis('off')
    plt.title('input')
    plt.imshow(plt_img)

    for idx_sim, j in enumerate(order_euclidean):
        if idx_sim>=TOP_NUM:
            break
        similar_img = images[j]
        plt.subplot(TEST_NUM, TOP_NUM+1, i * (TOP_NUM+1) + idx_sim + 2)
        plt.axis('off')
        plt.title('{:.2}'.format(distance_euclidean[idx_img,j]))
        plt.imshow(similar_img)

# 测试余弦距离
plt.figure()
for i,(idx_img,plt_img) in enumerate(test_images):
    order_cosin = np.argsort(distance_cosin[idx_img])[::-1]

    plt.subplot(TEST_NUM, TOP_NUM+1, i * (TOP_NUM+1) + 1)
    plt.axis('off')
    plt.title('input')
    plt.imshow(plt_img)

    for idx_sim, j in enumerate(order_cosin):
        if idx_sim>=TOP_NUM:
            break
        similar_img = images[j]
        plt.subplot(TEST_NUM, TOP_NUM+1,i* (TOP_NUM+1) + idx_sim + 2)
        plt.title('input')
        plt.axis('off')
        plt.title('{:.2}'.format(distance_cosin[idx_img,j]))
        plt.imshow(similar_img)

结果

结果展示的第一列是输入的图片，显然和其最相似的是它自己，然后是同类

Fig.1 欧氏距离：

Fig.2 余弦距离：