Implement Evaluation Metrics

repository/DBClient.py

@@ -14,3 +14,10 @@ def get_db_client():
         }
     )
     return db
+
+def get_db_client_for_relevance():
+    client = chromadb.PersistentClient(path=f"/data/dbs/{get_user_storage_id()}")
+    db = client.get_or_create_collection(
+        name = "relevance",
+    )
+    return db

repository/Search_Image.py

@@ -6,7 +6,7 @@ import numpy as np
 from io import BytesIO
 import time
 import bm25s
-from repository.DBClient import get_db_client
+from repository.DBClient import get_db_client, get_db_client_for_relevance
 from repository.Embedding_model import load_embedding_model
 import torch
 import torch.nn.functional as F
@@ -22,6 +22,7 @@ with st.sidebar:
 # show all images from chromadb in batch of 50
 try:
     db = get_db_client()
+    relevance_db = get_db_client_for_relevance()
 except Exception as e:
     st.warning("No images found. Please upload some images first.")
     st.stop()
@@ -62,8 +63,10 @@ if not search_query:
                     st.image(image, width="stretch")
                     if st.button("Delete", key=metadata["id"]):
                         db.delete(ids=[metadata["id"]])
+                        relevance_db.delete(ids=[metadata["id"]])
                         if os.path.exists(metadata["filepath"]):
                             os.remove(metadata["filepath"])
+                        
 
 else:
     start_time = time.time()
@@ -146,7 +149,8 @@ else:
         }
 
     best_metadata = dict(sorted(best_metadata.items(), key=lambda item: item[1]["combined_score"], reverse=True)[:k])
-
+    if "relevance_score" not in st.session_state:
+        st.session_state["relevance_score"] = [0] * k
     for filename, data in best_metadata.items():
         metadata = data["metadata"]
 
@@ -167,9 +171,37 @@ else:
                 f"Combined: {data['combined_score']:.2f}"
             )
         )
-        if st.button("Delete", key=f"delete_{metadata['id']}"):
-            db.delete(ids=[metadata["id"]])
-            if os.path.exists(metadata["filepath"]):
-                os.remove(metadata["filepath"])
-
+        columns = st.columns(3)
+        i = 1
+        with columns[0]:
+            if st.button("Delete", key=f"delete_{metadata['id']}"):
+                db.delete(ids=[metadata["id"]])
+                if os.path.exists(metadata["filepath"]):
+                    os.remove(metadata["filepath"])
+        with columns[1]:
+            if st.checkbox("Relevant 👍", key=f"relevant_{metadata['id']}"):
+                relevance_db.add(
+                    ids=[metadata["id"]],
+                    metadatas=[{**metadata, "relevance": 1, "query": search_query}]
+                )
+                st.session_state["relevance_score"][i - 1] = 1
+        with columns[2]:
+            if st.checkbox("Irrelevant 👎", key=f"irrelevant_{metadata['id']}"):
+                relevance_db.add(
+                    ids=[metadata["id"]],
+                    metadatas=[{**metadata, "relevance": 0, "query": search_query}]
+                )
+                st.session_state["relevance_score"][i - 1] = 0
+        i += 1
     st.write(f"Search took {time.time() - start_time:.2f} seconds")
+    st.info(f"You can provide relevance feedback by marking images as relevant or irrelevant. Ensure you have atleast {k} relevant images to evaluate the NDCG@{k} score.")
+    with st.button(f"Evaluate NDCG@{k}"):
+        ndcg_score = 0.0
+        if "relevance_score" in st.session_state:
+            relevance_scores = st.session_state["relevance_score"]
+            ideal_relevance_scores = sorted(relevance_scores, reverse=True)
+            dcg = sum((2 ** rel - 1) / math.log2(idx + 2) for idx, rel in enumerate(relevance_scores))
+            idcg = sum((2 ** rel - 1) / math.log2(idx + 2) for idx, rel in enumerate(ideal_relevance_scores))
+            ndcg_score = dcg / idcg if idcg > 0 else 0.0
+        st.title(f"NDCG@{k}: {ndcg_score:.4f}")
+        st.warning("Please note that the NDCG score is based on the relevance feedback you have provided. If you haven't marked any images as relevant or irrelevant, the score will be inaccurate.")