Spaces:
Sleeping
Sleeping
File size: 3,086 Bytes
b4f7e81 |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
import gradio as gr
from facenet_pytorch import InceptionResnetV1
import torch.nn as nn
import torchvision.transforms as tf
import numpy as np
import torch
import faiss
import h5py
import os
import random
from PIL import Image
import matplotlib.cm as cm
import matplotlib as mpl
img_names = []
with open('list_eval_partition.txt', 'r') as f:
for line in f:
img_name, dtype = line.rstrip().split(' ')
img_names.append(img_name)
# For a model pretrained on VGGFace2
print('Loading model weights ........')
class SiameseModel(nn.Module):
def __init__(self):
super().__init__()
self.backbone = InceptionResnetV1(pretrained='vggface2')
def forward(self, x):
x = self.backbone(x)
x = torch.nn.functional.normalize(x, dim=1)
return x
model = SiameseModel()
model.load_state_dict(torch.load('model.pt', map_location=torch.device('cpu')))
model.eval()
# Make FAISS index
print('Make index .............')
index = faiss.IndexFlatL2(512)
hf = h5py.File('face_vecs_full.h5', 'r')
for key in hf.keys():
vec = np.array(hf.get(key))
index.add(vec)
hf.close()
# Function to search image
def image_search(image, k=5):
transform = tf.Compose([
tf.Resize((160, 160)),
tf.ToTensor()
])
query_img = transform(image)
query_img = torch.unsqueeze(query_img, 0)
model.eval()
query_vec = model(query_img).detach().numpy()
D, I = index.search(query_vec, k=k)
retrieval_imgs = []
FOLDER = 'img_align_celeba'
for idx in I[0]:
img_file_name = img_names[idx]
path = os.path.join(FOLDER, img_file_name)
image = Image.open(path)
retrieval_imgs.append((image, ''))
return retrieval_imgs
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
gr.Markdown('''
# Face Image Retrieval with Content-based Retrieval Image (CBIR) & Saliency Map
--------
''')
with gr.Row():
with gr.Column():
image = gr.Image(type='pil', scale=1)
slider = gr.Slider(1, 10, value=5, step=1, label='Number of retrieval image')
with gr.Row():
btn = gr.Button('Search')
clear_btn = gr.ClearButton()
gallery = gr.Gallery(label='Retrieval Images', columns=[5], show_label=True, scale=2)
img_dir = './img_align_celeba'
examples = random.choices(img_names, k=6)
examples = [os.path.join(img_dir, ex) for ex in examples]
examples = [Image.open(img) for img in examples]
with gr.Row():
gr.Examples(
examples = examples,
inputs = image
)
btn.click(image_search,
inputs= [image, slider],
outputs= [gallery])
def clear_image():
return None
clear_btn.click(
fn = clear_image,
inputs = [],
outputs = [image]
)
def parse_args():
if __name__ == "__main__":
demo.launch()
|