基于VGG16神经网络实现以图搜图

思路

· 预先准备一份图片库，并对其中数据进行批处理操作，使用VGG16卷积神经网络提取图像的512维卷积特征，刷入数据库（ClickHouse）记录；

· 上传目标图像进行识图，同样使用VGG16提取目标图像特征，使用CK数据库距离函数进行匹配，高于阈值即可返回识图结果

神经网络

# -*- coding: utf-8 -*-
# @Author : tianL.R
# @Email : rtl1312@163.com
# @Time : 2023.11.26
import time

import numpy as np
from PIL import Image
from keras.applications.vgg16 import VGG16
from keras.applications.vgg16 import preprocess_input
from keras.preprocessing import image
from numpy import linalg


class VGG16Net:
    def __init__(self):
        self.input_shape = (224, 224, 3)
        self.weight = 'imagenet'
        self.pooling = 'max'
        self.model_vgg = VGG16(weights=self.weight,
                               input_shape=(self.input_shape[0], self.input_shape[1], self.input_shape[2],),
                               pooling=self.pooling,
                               include_top=False)
        self.model_vgg.predict(np.zeros((1, 224, 224, 3)))

    def detection(self, img_path):
        """
        提取VGG16最后一层卷积特征
        """
        # img = image.load_img(img_path, target_size=(self.input_shape[0], self.input_shape[1]))
        img = img_path.resize((self.input_shape[0], self.input_shape[1]))
        img = image.img_to_array(img)
        img = np.expand_dims(img, axis=0)
        img = preprocess_input(img)
        feat = self.model_vgg.predict(img)
        norm_feat = feat[0] / linalg.norm(feat[0])
        return norm_feat.tolist()


if __name__ == '__main__':
    img1 = '333.jpg'
    img2 = '555.jpg'
    img1 = Image.open(img1)
    img2 = Image.open(img2)

    vgg = VGG16Net()
    queryVec1 = np.array(vgg.detection(img1))
    queryVec2 = np.array(vgg.detection(img2))
    scores = np.dot(queryVec1, queryVec2)
    score2 = queryVec1.dot(queryVec2) / (np.linalg.norm(queryVec1) * np.linalg.norm(queryVec2))
    print(scores)
    print(score2)