added faiss

app.py

@@ -5,28 +5,42 @@ from setup import *
 
 from PIL import Image
 
+thumbnail_width = 300
 def show_result(search_request, 
                 search_result, 
                 img_dir,
                 container) :
 
-    # lorax =  Image.open('img/Lorax.jpg')
-    # print(lorax.width, lorax.height)
-    # st.image(lorax, width = 250)
-
     container.header("\"" +search_request+ "\" reminds me of :")
     i = 0
-    for _ in range(0, 2):
-        for col in container.columns(2) :
+    for _ in range(0, 3):
+        for col in container.columns(2):
+            if i >= len(search_result):
+                break
             image_name, comment, score = search_result[i]
-            col.image(img_dir + image_name, width = 300)
 
+            # Загрузка изображения
+            image = Image.open(img_dir + image_name)
+
+            # Выравнивание изображения по ширине
+            image_width, image_height = image.size
+            aspect_ratio = thumbnail_width / image_width
+            new_height = int(image_height * aspect_ratio)
+            resized_image = image.resize((thumbnail_width, new_height), Image.ANTIALIAS)
+
+            # Добавление подписи
             if score != '' :
-                sim_score = f"{float(100 * score):.3f}"
-                col.header(sim_score + "  " +comment)
+                 sim_score = f"{float(100 * score):.2f}"
+                 sim='similarity='+sim_score + "%"
+                 col.markdown(comment)
+                 col.markdown(f'<p style="font-size: 10px;">{sim}</p>', unsafe_allow_html=True)
             else :
-                col.header(comment)
+                # Вывод изображения в контейнер
+                col.markdown(comment)
+
+            col.image(resized_image, width=thumbnail_width)
             i = i + 1
+
     return
 
 def show_landing() :
@@ -52,6 +66,13 @@ def show_landing() :
                     search_result,
                     IMAGE_DIR+'/',
                     results)
+
+    if action.button('Find Relsease 3!') and os.path.exists(IMAGE_DIR) :
+        search_result = search3(search_request)
+        show_result(search_request,
+                    search_result,
+                    IMAGE_DIR+'/',
+                    results)
     return
 
 

dataframe.py

@@ -1,12 +1,11 @@
 import pandas as pd
 import numpy as np
 
-def get_image_data() :
+def get_image_data(csv_file) :
 
-    # flickr = pd.read_csv('data/results.csv', sep='|')
-    image_data_df = pd.read_csv ('data/output2.csv')
+    image_data_df = pd.read_csv (csv_file)
 
     image_data_df['text_embeddings'] = image_data_df['text_embeddings'].apply(lambda x: np.fromstring(x[2:-2], sep=' ')).values
     image_data_df['text_embeddings'] = image_data_df['text_embeddings'].apply(lambda x: np.reshape(x, (1, -1)))
 
-    return image_data_df
+    return image_data_df

main.py

@@ -38,10 +38,29 @@ def search2(search_prompt : str) :
     # Get model, processor & tokenizer
     model, tokenizer = get_model_info(model_ID, device)
 
-    image_data_df = get_image_data()
+    image_data_df = get_image_data('data/output2.csv')
 
     return get_top_N_images(search_prompt,
                             data = image_data_df,
                             model=model, tokenizer=tokenizer,
                             device = device,
-                            top_K=4)
+                            top_K=4)
+
+def search3(search_prompt : str) :
+
+    # Set the device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Define the model ID
+    model_ID = "openai/clip-vit-base-patch32"
+
+    # Get model, processor & tokenizer
+    model, tokenizer = get_model_info(model_ID, device)
+
+    image_data_df = get_image_data('data/output2.csv')
+
+    return faiss_get_top_N_images(search_prompt,
+                                  data = image_data_df,
+                                  model=model, tokenizer=tokenizer,
+                                  device = device,
+                                  top_K=4)

model.py

@@ -1,5 +1,6 @@
-from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer
+from transformers import CLIPModel, CLIPTokenizer
 from sklearn.metrics.pairwise import cosine_similarity
+import faiss
 
 from dataframe  import *
 
@@ -22,13 +23,13 @@ def get_single_text_embedding(text, model, tokenizer, device):
 
     return embedding_as_np
 
-def get_item_data(result, index) :
+def get_item_data(result, index, measure_column) :
 
     img_name = str(result['image_name'][index])
 
     # TODO: add code to get the original comment
     comment = str(result['comment'][index])
-    cos_sim = result['cos_sim'][index]
+    cos_sim = result[measure_column][index]
 
     return (img_name, comment, cos_sim)
 
@@ -36,29 +37,81 @@ def get_top_N_images(query,
                      data,
                      model, tokenizer,
                      device,
-                     top_K=4, 
-                     search_criterion="text"):
-   # Text to image Search
-   if (search_criterion.lower() == "text"):
-     query_vect = get_single_text_embedding(query, model, tokenizer, device)
-#    # Image to image Search
-#    else:
-#      query_vect = get_single_image_embedding(query)
-
-   # Relevant columns
-   revevant_cols = ["comment", "image_name", "cos_sim"]
-
-   # Run similarity Search
-   data["cos_sim"] = data["text_embeddings"].apply(lambda x: cosine_similarity(query_vect, x))# line 17
-   data["cos_sim"] = data["cos_sim"].apply(lambda x: x[0][0])
-
-   data_sorted = data.sort_values(by='cos_sim', ascending=False)
-   non_repeated_images = ~data_sorted["image_name"].duplicated()
-   most_similar_articles = data_sorted[non_repeated_images].head(top_K)
-
-   """
-   Retrieve top_K (4 is default value) articles similar to the query
-   """
-
-   result_df = most_similar_articles[revevant_cols].reset_index()
-   return [get_item_data(result_df, i) for i in range(len(result_df))]
+                     top_K=4) :
+
+    query_vect = get_single_text_embedding(query, 
+                                            model, tokenizer, 
+                                            device)
+
+    # Relevant columns
+    relevant_cols = ["comment", "image_name", "cos_sim"]
+
+    # Run similarity Search
+    data["cos_sim"] = data["text_embeddings"].apply(lambda x: cosine_similarity(query_vect, x))# line 17
+    data["cos_sim"] = data["cos_sim"].apply(lambda x: x[0][0])
+
+    data_sorted = data.sort_values(by='cos_sim', ascending=False)
+    non_repeated_images = ~data_sorted["image_name"].duplicated()
+    most_similar_articles = data_sorted[non_repeated_images].head(top_K)
+
+    """
+    Retrieve top_K (4 is default value) articles similar to the query
+    """
+
+    result_df = most_similar_articles[relevant_cols].reset_index()
+
+    return [get_item_data(result_df, i, 'cos_sim') for i in range(len(result_df))]
+
+
+###### with faiss ###########
+
+import faiss
+import numpy as np
+
+def faiss_add_index_cos(df, column):
+
+    # Get the embeddings from the specified column
+    embeddings = np.vstack(df[column].values).astype(np.float32)  # Convert to float32
+          
+    # Create an index
+    index = faiss.IndexFlatIP(embeddings.shape[1])
+    print("<<<<faiss_ after normalize") 
+    faiss.normalize_L2(embeddings)
+    print("<<<<faiss_ after normalize") 
+
+    index.train(embeddings)
+    print("<<<<faiss_ after index.train")
+
+    # Add the embeddings to the index
+    index.add(embeddings)
+    print("<<<<faiss_add")
+
+    # Return the index
+    return index
+
+
+def faiss_get_top_N_images(query,
+                           data,
+                           model, tokenizer,
+                           device,
+                           top_K=4) :
+
+    query_vect = get_single_text_embedding(query, 
+                                          model, tokenizer, 
+                                          device)
+    # Relevant columns
+    relevant_cols = ["comment", "image_name", "similarity"]
+
+    #faiss search with cos similarity
+    index = faiss_add_index_cos(data, column="text_embeddings")
+
+    faiss.normalize_L2(query_vect)
+    D, I = index.search(query_vect, len(data))
+
+    data_sorted = data.iloc[I.flatten()]
+
+    non_repeated_images = ~data_sorted["image_name"].duplicated()
+    most_similar_articles = data_sorted[non_repeated_images].head(top_K)
+
+    result_df = most_similar_articles[relevant_cols].reset_index(), D.reshape(-1,1)[:top_K]
+    return [get_item_data(result_df, i, 'similarity') for i in range(len(result_df))]