Upload 4 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model2.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,59 @@
+import gradio as gr
+import numpy as np
+from transformers import AutoFeatureExtractor, AutoModel
+from datasets import load_dataset
+from PIL import Image, ImageDraw
+import os
+
+    
+# 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
+
+# Load dataset
+dataset_with_embeddings = load_dataset("tonyassi/vogue-runway-top15-512px-nobg-embeddings2", split="train")
+dataset_with_embeddings.add_faiss_index(column='embeddings')
+
+
+def get_neighbors(query_image, top_k=10):
+    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(image_dict):
+
+    # Open query image
+    query_image = Image.open(image_dict['composite']).convert(mode='RGB')
+
+    # Get similar image
+    scores, retrieved_examples = get_neighbors(query_image)
+
+    #final_md = ""
+    
+    # Create result diction for gr.Gallery
+    result = []
+    for i in range(len(retrieved_examples["image"])):
+        id = retrieved_examples["label"][i]
+        print('id', id)
+        label = dataset_with_embeddings.features["label"].names[id]
+        print('label', label)
+        result.append((retrieved_examples["image"][i], label))
+    
+    return result, query_image
+
+iface = gr.Interface(fn=search,
+                     title='Sketch to Fashion Collection',
+                     description="""
+                     Tony Assi
+                     """,
+                     inputs=gr.ImageEditor(label='Sketchpad' ,type='filepath', value={'background':'./model2.png', 'layers':None, 'composite':None}, sources=['upload'], transforms=[]), 
+                     outputs=[gr.Gallery(label='Similar', object_fit='contain', height=900), gr.Image()],
+                     #examples=[[{'background':'./images/goth.jpg', 'layers':None, 'composite':'./images/goth.jpg'}],[{'background':'./images/pink.jpg', 'layers':None, 'composite':'./images/pink.jpg'}], [{'background':'./images/boot.jpg', 'layers':None, 'composite':'./images/boot.jpg'}]],
+                     theme = gr.themes.Base(primary_hue="teal",secondary_hue="teal",neutral_hue="slate"),)
+iface.launch()

requirements.txt

@@ -0,0 +1,5 @@
+transformers
+torch
+datasets
+faiss-cpu
+https://gradio-builds.s3.amazonaws.com/f3e3c5c02f1069c1180004a46909309544f42523/gradio-4.16.0-py3-none-any.whl