add nearest neighbor

.gitattributes

@@ -26,3 +26,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.faiss filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,5 +1,6 @@
+import re
 import streamlit as st # HF spaces at v1.2.0
-from demo import load_model,generate
+from demo import load_model,generate,get_dataset,embed
 
 # TODOs
 # Add markdown short readme project intro
@@ -21,21 +22,46 @@ def load_model_intocache(model_name):
 
     return gan
 
+@st.experimental_singleton
+def load_dataset():
+    dataset=get_dataset()
+    return dataset
+
 model_name='ceyda/butterfly_cropped_uniq1K_512'
 model=load_model_intocache(model_name)
+dataset=load_dataset()
 
 st.write(f"Model {model_name} is loaded")
 st.write(f"Latent dimension: {model.latent_dim}, Image size:{model.image_size}")
 
-run=st.button("Generate")
-if run:
+if 'ims' not in st.session_state:
+    st.session_state['ims'] = None
+
+ims=st.session_state["ims"]
+batch_size=4 #generate 4 butterflies 
+def run():
     with st.spinner("Generating..."):
-    
-        batch_size=4 #generate 4 butterflies 
         ims=generate(model,batch_size)
+        st.session_state['ims'] = ims
 
-        cols=st.columns(batch_size)
-        for i,im in enumerate(ims):
-            cols[i].image(im)
+runb=st.button("Generate", on_click=run)
+if ims is not None:
+    cols=st.columns(batch_size)
+    picks=[False]*batch_size
+    for i,im in enumerate(ims):
+        cols[i].image(im)
+        picks[i]=cols[i].button("Find Nearest",key="pick_"+str(i))
+        # if picks[i]:
+        #     scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(im), k=5)
+        #     for r in retrieved_examples["image"]:
+        #         st.image(r)
+    
+if any(picks):
+    # st.write("Nearest butterflies:")
+    for i,pick in enumerate(picks):
+        if pick:
+            scores, retrieved_examples=dataset.get_nearest_examples('beit_embeddings', embed(ims[i]), k=5)
+            for r in retrieved_examples["image"]:
+                cols[i].image(r)
 
 

beit_index.faiss

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

demo.py

@@ -7,15 +7,34 @@ def get_train_data(dataset_name="ceyda/smithsonian_butterflies_transparent_cropp
     dataset=dataset.sort("sim_score")
     score_thresh = dataset["train"][data_limit]['sim_score']
     dataset = dataset.filter(lambda x: x['sim_score'] < score_thresh)
-    
-    dataset = dataset.map(lambda x: x.convert("RGB"))
+    dataset = dataset.map(lambda x: {'image' : x['image'].convert("RGB")})
     return dataset["train"]
+
+from transformers import BeitFeatureExtractor, BeitForImageClassification
+feature_extractor = BeitFeatureExtractor.from_pretrained('microsoft/beit-base-patch16-224')
+model = BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224')
+def embed(images):
+    inputs = feature_extractor(images=images, return_tensors="pt")
+    outputs = model(**inputs,output_hidden_states= True)
+    last_hidden=outputs.hidden_states[-1]
+    pooler=model.base_model.pooler
+    final_emb=pooler(last_hidden).detach().numpy()
+    return final_emb    
     
-    
+def build_index():    
+    dataset=get_train_data()
+    ds_with_embeddings = dataset.map(lambda x: {"beit_embeddings":embed(x["image"])},batched=True,batch_size=20)
+    ds_with_embeddings.add_faiss_index(column='beit_embeddings')
+    ds_with_embeddings.save_faiss_index('beit_embeddings', 'beit_index.faiss')
+
+def get_dataset():
+    dataset=get_train_data()
+    dataset.load_faiss_index('beit_embeddings', 'beit_index.faiss')
+    return dataset
 
 def load_model(model_name='ceyda/butterfly_cropped_uniq1K_512'):
     gan = LightweightGAN.from_pretrained(model_name)
-    gan.eval();
+    gan.eval()
     return gan
     
 def generate(gan,batch_size=1):