-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathapp.py
95 lines (75 loc) · 2.74 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from PIL import Image
from similarity_retrieval import database
from similarity_retrieval.model import LatentModel, get_pretrained_model
from similarity_retrieval.database import download_fashion_mnist, DEFAULT_PATH
def load_image(image_file):
img = Image.open(image_file)
return np.array(img)
def plot_images(images):
plt.figure(figsize=(20, 10))
columns = 5
for (i, image) in enumerate(images):
ax = plt.subplot(len(images) // columns + 1, columns, i + 1)
ax.set_title("Similar Image " + str(i))
plt.imshow(image)
plt.axis("off")
st.pyplot(plt)
def visualize_query_results(
latent_model, query_image, training_images, no_of_results, show_plots=False
):
ids = latent_model.query(query_image)
candidates = []
for idx, id in enumerate(ids):
if idx == no_of_results:
break
candidates.append(training_images[int(id)])
# candidates.insert(0, query_image)
plot_images(candidates)
if __name__ == "__main__":
st.header("Image Similarity Retrieval")
(x_train, y_train), (x_test, y_test) = download_fashion_mnist(samples=6000)
training_files = zip(x_train, y_train)
embedding_model = get_pretrained_model()
latent_model = LatentModel(embedding_model)
latent_model.train(training_files)
val_images = []
val_labels = []
for image, label in zip(x_test[:50], y_test[:50]):
val_images.append(image)
val_labels.append(label)
val_images = np.array(val_images)
val_labels = np.array(val_labels)
no_of_results = 10
st.subheader("Image Query")
query_type = st.radio("Query Type", ("Random", "Image Upload"))
if query_type == "Image Upload":
image_file = st.file_uploader(
"Upload Query Image",
type=["png", "jpg", "jpeg"]
)
if image_file is not None:
st.subheader("Uploaded Image")
# To View Uploaded Image
st.image(load_image(image_file), width=75)
query_image = load_image(image_file)
else:
query_image = None
else:
st.subheader("Random Image")
st.caption("Randomly selected image from the test set")
idx = np.random.choice(len(val_images) - 1)
query_image = val_images[idx]
st.image(query_image, width=75)
if query_image is not None:
st.subheader("Filter")
no_of_results = st.slider(
"Number of Similar Images to show", min_value=10, max_value=len(val_images)
)
st.subheader("Results")
visualize_query_results(
latent_model, query_image, x_train, no_of_results, show_plots=True
)