init

app.py

@@ -0,0 +1,117 @@
+import gradio as gr
+
+import numpy as np
+import h5py
+import pandas as pd
+import pickle
+
+from skimage.util import random_noise
+import time
+
+import tensorflow as tf
+
+from model import VGGNet
+
+def query_normal(query_img: np.ndarray):
+    query_feature = model.extract_feature(query_img, verbose=0)
+    scores = np.dot(query_feature, features.T)
+    rank_ID = np.argsort(scores)[::-1]
+    rank_score = scores[rank_ID]
+
+    imlist = [image_ids.astype(str)[index] for i, index in enumerate(rank_ID[0:10])]
+
+    return imlist, rank_score[0:10].tolist(), len(scores)
+
+def query_kmeans(query_img: np.ndarray):
+    query_feature = model.extract_feature(query_img, verbose=0)
+
+    cluster = kmeans.predict(query_feature.reshape(1, -1))
+    cluster = cluster[0]
+
+    df = pd.DataFrame({'image_id': image_ids.astype(str), 'cluster_id': kmeans.labels_})
+    df = df[df['cluster_id'] == cluster]
+
+    query = df[df['cluster_id'] == cluster].index
+
+    query_feature = model.extract_feature(query_img)
+    scores = np.dot(query_feature, features[query].T)
+    rank_ID = np.argsort(scores)[::-1]
+    rank_score = scores[rank_ID]
+
+    imlist = [image_ids.astype(str)[index] for i, index in enumerate(query[rank_ID[0:10]])]
+
+    return imlist, rank_score[0:10].tolist(), len(scores), cluster
+
+def query(image, mode, noise, noise_seed, mean, var, amount, salt_vs_pepper):
+    if image == None or noise_seed == None:
+        return None, None, None, None, None
+    
+    query_img = tf.keras.utils.load_img(image, target_size=(model.input_shape[0], model.input_shape[1]))
+    query_img = tf.keras.utils.img_to_array(query_img).astype(int)
+
+    if noise == 'none':
+        pass
+    elif noise == 'gaussian' or noise =='speckle':
+        query_img = random_noise(query_img / 255, mode=noise, rng=int(noise_seed), mean=mean, var=var, clip=True)
+        query_img = np.array(query_img * 255, dtype=np.uint8)
+    elif noise == 'localvar':
+        query_img = random_noise(query_img / 255, mode=noise, rng=int(noise_seed), clip=True)
+        query_img = np.array(query_img * 255, dtype=np.uint8)
+    elif noise == 'poisson':
+        query_img = random_noise(query_img / 255, mode=noise, rng=int(noise_seed), clip=True)
+        query_img = np.array(query_img * 255, dtype=np.uint8)
+    elif noise == 'salt' or noise == 'pepper':
+        query_img = random_noise(query_img / 255, mode=noise, rng=int(noise_seed), amount=amount, clip=True)
+        query_img = np.array(query_img * 255, dtype=np.uint8)
+    elif noise == 's&p':
+        query_img = random_noise(query_img / 255, mode=noise, rng=int(noise_seed), amount=amount, salt_vs_pepper=salt_vs_pepper, clip=True)
+        query_img = np.array(query_img * 255, dtype=np.uint8)
+
+    start = time.time()
+    if mode == 'normal':
+        results, scores, length = query_normal(query_img)
+    elif mode == 'kmeans':
+        results, scores, length, cluster = query_kmeans(query_img)
+    end = time.time()
+    query_time = end - start
+    query_time = round(query_time * 1000, 2)
+
+    return query_img, f'Query time: {query_time} ms', [(result, f'Score: {score}, file: {result}') for result, score in zip(results, scores)], length, f'{cluster}' if mode == 'kmeans' else 'None'
+
+if __name__ == '__main__':
+    model = VGGNet()
+
+    # Load dataset
+    datasets = h5py.File('features.h5', 'r')
+    features = datasets['features'][:]
+    image_ids = datasets['image_ids'][:]
+    datasets.close()
+
+    # Load kmeans model
+    with open('kmeans.pkl', 'rb') as f:
+        kmeans = pickle.load(f)
+
+    # Run web app
+    iface = gr.Interface(
+        title='An image search engine based on VGG16',
+        description='dataset: https://www.kaggle.com/datasets/hcfighting/tinyimagenet200. Error may occur when result showing because i\'d deleted some images in the dataset, please try another image to search.',
+        fn=query,
+        inputs=[
+            gr.Image(type='filepath'),
+            gr.Radio(['normal', 'kmeans'], value='normal', info='normal: search all images, kmeans: search images in the same cluster'),
+            gr.Radio(['none', 'gaussian', 'localvar', 'poisson', 'salt', 'pepper', 's&p', 'speckle'], value='none', info='noise type (default: none)'),
+            gr.Number(label='random_seed', value=0, info='random seed for noise (default: 0)'),
+            gr.Slider(label='mean', maximum=1, minimum=0, value=0, info='mean of the noise, works with gaussian and localvar (default: 0)'),
+            gr.Slider(label='var', maximum=1, minimum=0, value=0.01, info='variance of the noise, works with gaussian and localvar (default: 0.01)'),
+            gr.Slider(label='amount', maximum=1, minimum=0, value=0.05, info='proportion of image pixels to replace with noise, works with salt, pepper, and s&p (default: 0.05)'),
+            gr.Slider(label='salt_vs_pepper', maximum=1, minimum=0, value=0.5, info='proportion of s&p noise (default: 0.5)'),
+        ],
+        outputs=[
+            gr.Image(label='query image', type='numpy'),
+            gr.Label(label='query time'),
+            gr.Gallery(label='Top 10 similar images').style(columns=[5], rows=[2]),
+            gr.Label(label='count of images for searching'),
+            gr.Label(label='cluster id')
+        ],
+    )
+    iface.launch()

features.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:213bc8f61b17123b4f6977843eeef32d479c176b03c65fb6f9d404e8a704ac4a
+size 210702048

kmeans.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4634447f64ae7d12f0b0a04bf49a806a4bb2f3b13d308747deaa9942f136b153
+size 421136

model.py

@@ -0,0 +1,34 @@
+import tensorflow as tf
+
+import numpy as np
+from numpy import linalg as LA
+
+from keras.applications.vgg16 import VGG16
+from keras.applications.vgg16 import preprocess_input
+
+class VGGNet:
+    def __init__(self):
+        self.input_shape = (224, 224, 3)
+        self.weight = 'imagenet'
+        self.pooling = 'max'
+        self.model = 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.predict(np.zeros((1, 224, 224, 3)))
+
+    def extract_feature(self, img_path: any, verbose='auto'):
+        if type(img_path) == str or type(img_path) == np.str_:
+            img = tf.keras.utils.load_img(img_path, target_size=(self.input_shape[0], self.input_shape[1]))
+            img = tf.keras.utils.img_to_array(img)
+            img = np.expand_dims(img, axis=0)
+            img = preprocess_input(img)
+            feature = self.model.predict(img, verbose=verbose)
+            feature_normalized = feature[0] / LA.norm(feature[0])
+        elif type(img_path) == np.ndarray:
+            img = img_path
+            img = np.expand_dims(img, axis=0)
+            img = preprocess_input(img)
+            feature = self.model.predict(img, verbose=verbose)
+            feature_normalized = feature[0] / LA.norm(feature[0])
+        else:
+            raise TypeError('img_path must be str or np.ndarray')
+        
+        return feature_normalized

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+h5py
+scikit_learn
+scikit-image
+tensorflow