use corresponding embeddings for each model

app.py

@@ -196,13 +196,13 @@ def get_embeddings(input_image,model_name):
 
 def find_closest(input_image,model_name):
     embedding = get_embeddings(input_image,model_name)
-    classes, paths = get_images(embedding)
+    classes, paths = get_images(embedding,model_name)
     #outputs = classes+paths
     return classes,paths
 
 def generate_diagram_closest(input_image,model_name,top_k):
     embedding = get_embeddings(input_image,model_name)
-    diagram_path = get_diagram(embedding,top_k)
+    diagram_path = get_diagram(embedding,top_k,model_name)
     return diagram_path
 
 def explain_image(input_image,model_name,explain_method,nb_samples):

closest_sample.py

@@ -9,9 +9,6 @@ import matplotlib.pyplot as plt
 from collections import Counter
 
 
-pca_fossils = pk.load(open('pca_fossils_142_resnet.pkl','rb'))
-pca_leaves = pk.load(open('pca_leaves_142_resnet.pkl','rb'))
-
 if not os.path.exists('dataset'):
   REPO_ID='Serrelab/Fossils'
   token = os.environ.get('READ_TOKEN')
@@ -20,8 +17,6 @@ if not os.path.exists('dataset'):
      print("warning! A read token in env variables is needed for authentication.")
   snapshot_download(repo_id=REPO_ID, token=token,repo_type='dataset',local_dir='dataset')
 
-embedding_fossils = np.load('dataset/embedding_leaves_142_finer.npy')
-#embedding_leaves = np.load('embedding_leaves.npy')
 
 fossils_pd= pd.read_csv('fossils_paths.csv')
 
@@ -57,8 +52,20 @@ def download_public_image(url, destination_path):
     else:
         print(f"Failed to download image from bucket. Status code: {response.status_code}")
 
-def get_images(embedding):
-    
+def get_images(embedding,model_name):
+
+    if model_name in ['Rock 170','Mummified 170']:
+        pca_fossils = pk.load(open('pca_fossils_170_finer.pkl','rb'))
+        pca_leaves = pk.load(open('pca_leaves_170_finer.pkl','rb'))
+        embedding_fossils = np.load('dataset/embedding_fossils_170_finer.npy')
+        #embedding_leaves = np.load('embedding_leaves.npy')
+    elif model_name in ['Fossils 142']:
+        pca_fossils = pk.load(open('pca_fossils_142_resnet.pkl','rb'))
+        pca_leaves = pk.load(open('pca_leaves_142_resnet.pkl','rb'))
+        embedding_fossils = np.load('dataset/embedding_leaves_142_finer.npy')
+        #embedding_leaves = np.load('embedding_leaves.npy')
+    else:
+        print(f'{model_name} not recognized')
     #pca_embedding_fossils = pca_fossils.transform(embedding_fossils[:,-1])
     
     pca_d =pca_distance(pca_fossils,embedding,embedding_fossils,top_k=5)
@@ -93,7 +100,20 @@ def get_images(embedding):
 
     return classes, local_paths
 
-def get_diagram(embedding,top_k):
+def get_diagram(embedding,top_k,model_name):
+
+    if model_name in ['Rock 170','Mummified 170']:
+        pca_fossils = pk.load(open('pca_fossils_170_finer.pkl','rb'))
+        pca_leaves = pk.load(open('pca_leaves_170_finer.pkl','rb'))
+        embedding_fossils = np.load('dataset/embedding_fossils_170_finer.npy')
+        #embedding_leaves = np.load('embedding_leaves.npy')
+    elif model_name in ['Fossils 142']:
+        pca_fossils = pk.load(open('pca_fossils_142_resnet.pkl','rb'))
+        pca_leaves = pk.load(open('pca_leaves_142_resnet.pkl','rb'))
+        embedding_fossils = np.load('dataset/embedding_leaves_142_finer.npy')
+        #embedding_leaves = np.load('embedding_leaves.npy')
+    else:
+        print(f'{model_name} not recognized')
     
     #pca_embedding_fossils = pca_fossils.transform(embedding_fossils[:,-1])
     

test.py

@@ -23,7 +23,7 @@
 import numpy as np
 
 # Load the .npy file
-embedding = np.load('embedding.npy')
+embedding = np.load('dataset/embedding_leaves_142_finer.npy')
 
 # Check the shape of the array
 print(embedding.shape)