Upload 2 files

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+import numpy as np
+from transformers import AutoFeatureExtractor, AutoModel
+from datasets import load_dataset
+from PIL import Image, ImageDraw
+
+
+    
+
+# Load model for computing embeddings of the candidate images
+print('Load model for computing embeddings of the candidate images')
+model_ckpt = "google/vit-base-patch16-224"
+extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
+model = AutoModel.from_pretrained(model_ckpt)
+hidden_dim = model.config.hidden_size
+
+
+dataset_with_embeddings = load_dataset("LucyintheSky/24-1-8-ds-embeddings", split="train")
+dataset_with_embeddings.add_faiss_index(column='embeddings')
+
+
+def get_neighbors(query_image, top_k=8):
+    qi_embedding = model(**extractor(query_image, return_tensors="pt"))
+    qi_embedding = qi_embedding.last_hidden_state[:, 0].detach().numpy().squeeze()
+    scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples('embeddings', qi_embedding, k=top_k)
+    return scores, retrieved_examples
+
+    
+
+def search(img):
+    
+    query_image = Image.open(img).convert(mode='RGB')
+
+    #query_image.thumbnail((1000,1000))
+
+
+    #query_image = query_image.resize((160,160))
+
+    print('search')
+    scores, retrieved_examples = get_neighbors(query_image)
+
+    print('return example')
+
+    result = []
+    for i in range(len(retrieved_examples["image"])):
+        id = str(retrieved_examples["text"][i]) + ' ' + str(scores[i])
+        print('id', id)
+        #label = dataset_with_embeddings.features["label"].names[id]
+        #print('label', label)
+        result.append((retrieved_examples["image"][i], id))
+    
+    return result, query_image
+
+iface = gr.Interface(fn=search,
+                     title='Celebrity Look-a-Like',
+                     inputs=gr.Image(type='filepath', label='Your Photo'), 
+                     outputs=[gr.Gallery(label='Similar', object_fit='contain'), gr.Image(label='Face')],
+                     #examples=[['./images/tony.jpg'],['./images/jessica.jpg'],['./images/scarlett.jpg'],['./images/christian.jpg']],
+                     theme = gr.themes.Base(primary_hue="teal",secondary_hue="teal",neutral_hue="slate"),)
+iface.launch()

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+torch
+datasets
+faiss-cpu