Update app.py

app.py

@@ -5,10 +5,10 @@ import random
 import numpy as np
 import time
 
-
-#ds = load_dataset("tonyassi/lucy4-embeddings", split='train')
+# Dataset
 ds = load_dataset("tonyassi/finesse1-embeddings", split='train')
-#ds = load_dataset("tonyassi/lucy5-embeddings", split='train')
+
+
 id_to_row = {row['id']: row for row in ds}
 remaining_ds = None
 preference_embedding = []
@@ -28,41 +28,13 @@ def get_random_images(dataset, num):
     
     return random_images, new_dataset
 
-"""
-def find_similar_images(dataset, num, embedding):
-    start_time = time.time()
-    # Find the most similar images in dataset
-    dataset.add_faiss_index(column='embeddings')
-    embedding = np.array(embedding)
-    scores, retrieved_examples = dataset.get_nearest_examples('embeddings', embedding, k=num)
-
-    print('time 2.1:', time.time()-start_time)
-
-    # Create a new dataset without these images
-    dataset.drop_index('embeddings')
-    print('time 2.2:', time.time()-start_time)
-    remaining_indices = [i for i in range(len(dataset)) if dataset[i]['id'] not in retrieved_examples['id']]
-    print('time 2.3:', time.time()-start_time)
-    new_dataset = dataset.select(remaining_indices)
-
-    print('time 2.4:', time.time()-start_time)
-    return retrieved_examples, new_dataset
-
-"""
-
 def find_similar_images(dataset, num, embedding):
-    start_time = time.time()
-    
     # Ensure FAISS index exists and search for similar images
-    #if not dataset.has_faiss_index('embeddings'):
     dataset.add_faiss_index(column='embeddings')
     scores, retrieved_examples = dataset.get_nearest_examples('embeddings', np.array(embedding), k=num)
     
-    print('time 2.1:', time.time()-start_time)
-    
     # Drop FAISS index after use to avoid re-indexing
     dataset.drop_index('embeddings')
-    print('time 2.2:', time.time()-start_time)
 
     # Extract all dataset IDs and use a set to find remaining indices
     dataset_ids = dataset['id']
@@ -70,17 +42,12 @@ def find_similar_images(dataset, num, embedding):
 
     # Use a list comprehension with enumerate for faster indexing
     remaining_indices = [i for i, id in enumerate(dataset_ids) if id not in retrieved_ids_set]
-    
-    print('time 2.3:', time.time()-start_time)
 
     # Create a new dataset without the retrieved images
     new_dataset = dataset.select(remaining_indices)
 
-    print('time 2.4:', time.time()-start_time)
     return retrieved_examples, new_dataset
 
-
-
 def average_embedding(embedding1, embedding2):
     embedding1 = np.array(embedding1)
     embedding2 = np.array(embedding2)
@@ -89,7 +56,6 @@ def average_embedding(embedding1, embedding2):
 ###################################################################################
 
 def load_images():
-    print('load_images()')
     print("ds", ds.num_rows)
 
     global remaining_ds
@@ -108,23 +74,15 @@ def load_images():
 
 
 def select_image(evt: gr.SelectData, gallery, preference_gallery):
-    start_time = time.time()
-
-    print('select_image()')
-
     global remaining_ds
     print("remaining_ds", remaining_ds.num_rows)
     
     # Selected image
     selected_id = int(evt.value['caption'])
-    print('ID', selected_id)
-    #selected_row = ds.filter(lambda row: row['id'] == selected_id)[0]
     selected_row = id_to_row[selected_id]
     selected_embedding = selected_row['embeddings']
     selected_image = selected_row['image']
 
-    print('time 1:', time.time()-start_time)
-
     # Update preference embedding
     global preference_embedding
     if len(preference_embedding) == 0:
@@ -132,18 +90,12 @@ def select_image(evt: gr.SelectData, gallery, preference_gallery):
     else: 
         preference_embedding = average_embedding(preference_embedding, selected_embedding)
 
-    print('time 2:', time.time()-start_time)
-
     # Find images which are most similar to the preference embedding
     simlar_images, remaining_ds = find_similar_images(remaining_ds, 5, preference_embedding)
 
-    print('time 3:', time.time()-start_time)
-
     # Create a list of tuples [(img1,caption1),(img2,caption2)...]
     result = list(zip(simlar_images['image'], [str(id) for id in simlar_images['id']]))
 
-    print('time 4:', time.time()-start_time)
-
     # Get random images
     rand_imgs, remaining_ds = get_random_images(remaining_ds, 5)
     # Create a list of tuples [(img1,caption1),(img2,caption2)...]
@@ -157,8 +109,6 @@ def select_image(evt: gr.SelectData, gallery, preference_gallery):
     else:
         final_preference_gallery = [selected_image] + preference_gallery
 
-    print('time 5:', time.time()-start_time)
-
     return gr.Gallery(value=final_result, selected_index=None), final_preference_gallery
 
 ###################################################################################