Add: Saving index to faiss file

src/app.py

@@ -23,13 +23,27 @@ def index_dataset(dataset):
 
     # Add index
     dataset_with_embeddings.add_faiss_index(column='color_embeddings')
+    dataset_with_embeddings.save_faiss_index('color_embeddings', 'color_index.faiss')
+
     dataset_with_embeddings.add_faiss_index(column='clip_embeddings')
+    ds_with_embeddings.save_faiss_index('clip_embeddings', 'clip_index.faiss')
 
 
     print(dataset_with_embeddings)
     return dataset_with_embeddings
 
-dataset_with_embeddings = index_dataset(candidate_subset)
+
+def check_index(ds):
+    index_path = "my_index.faiss"
+    if os.path.isfile('color_index.faiss') and os.path.isfile('clip_index.faiss'):
+        ds.load_faiss_index('color_embeddings', 'color_index.faiss')
+        return ds.load_faiss_index('clip_embeddings', 'clip_index.faiss')
+
+    else:
+        return index_dataset(ds)
+
+
+dataset_with_embeddings = check_index(candidate_subset)
 
 # Main function, to find similar images
 # TODO: implement different distance measures
@@ -66,35 +80,16 @@ def get_neighbors(query_image, selected_descriptor, top_k=5):
 
 # Define the Gradio Interface
 
-
-iface = gr.Interface(
-    fn=get_neighbors,
-    inputs=[
-        gr.Image(type="pil", label="Your Image"),
-        gr.CheckboxGroup(["Color Descriptor", "LBP", "CLIP"], label="Descriptor method?"),
-    ],
-    outputs=gr.Gallery(),
-    title="Image Similarity Gallery",
-    description="Upload an image and get similar images",
-    allow_flagging="never"
-)
-
-# Launch the Gradio interface
-iface.launch()
-
 with gr.Blocks() as demo:
     image_input = gr.Image(type="pil", label="Please upload an image")
     checkboxes_descr = gr.CheckboxGroup(["Color Descriptor", "LBP", "CLIP"], label="Please choose an descriptor"))
 
     btn = gr.Button(value="Submit")
     gallery_output = gr.Gallery()
+    btn.click(get_neighbors, inputs=[image_input, checkboxes_descr], outputs=[gallery_output])
 
-    btn.click(get_neighbors,
-              inputs=[image_input, checkboxes_descr],
-              outputs=[gallery_output]
-              )
     btn_index = gr.Button(value="Re-index Dataset")
-    btn.click(index_dataset)
+    btn_index.click(index_dataset)
 
 if __name__ == "__main__":
     demo.launch()