Update src/streamlit_app.py

src/streamlit_app.py

@@ -4,21 +4,45 @@ import clip
 from PIL import Image
 import os
 import numpy as np
+import chromadb
+from chromadb.utils import embedding_functions
 
 # Initialize session state
 if 'model' not in st.session_state:
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model, preprocess = clip.load("ViT-B/32", device=device)
+    # Set a custom cache directory for CLIP model weights
+    cache_dir = "./clip_cache"
+    os.makedirs(cache_dir, exist_ok=True)  # Create cache directory if it doesn't exist
+    try:
+        model, preprocess = clip.load("ViT-B/32", device=device, download_root=cache_dir)
+    except Exception as e:
+        st.error(f"Failed to load CLIP model: {e}")
+        st.stop()
     st.session_state.model = model
     st.session_state.preprocess = preprocess
     st.session_state.device = device
     st.session_state.demo_images = []
-    st.session_state.demo_encodings = []
     st.session_state.demo_image_paths = []
     st.session_state.user_images = []
-    st.session_state.user_encodings = []
 
-# Load demo images
+# Initialize ChromaDB client
+if 'chroma_client' not in st.session_state:
+    try:
+        st.session_state.chroma_client = chromadb.PersistentClient(path="./chroma_db")
+        # Create or get collections
+        st.session_state.demo_collection = st.session_state.chroma_client.get_or_create_collection(
+            name="demo_images",
+            metadata={"hnsw:space": "cosine"}  # Use cosine similarity
+        )
+        st.session_state.user_collection = st.session_state.chroma_client.get_or_create_collection(
+            name="user_images",
+            metadata={"hnsw:space": "cosine"}
+        )
+    except Exception as e:
+        st.error(f"Failed to initialize ChromaDB collections: {e}")
+        st.stop()
+
+# Load demo images into ChromaDB
 if not st.session_state.demo_images:
     demo_folder = "demo_images"
     if os.path.exists(demo_folder):
@@ -26,10 +50,31 @@ if not st.session_state.demo_images:
         if len(demo_image_paths) > 0:
             st.session_state.demo_image_paths = demo_image_paths
             st.session_state.demo_images = [Image.open(path) for path in demo_image_paths]
-            demo_preprocessed = [st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device) for img in st.session_state.demo_images]
-            with torch.no_grad():
-                demo_encodings = [st.session_state.model.encode_image(img) for img in demo_preprocessed]
-            st.session_state.demo_encodings = torch.cat(demo_encodings, dim=0)
+            
+            # Clear existing demo collection to avoid duplicates
+            st.session_state.demo_collection.delete(ids=[str(i) for i in range(len(demo_image_paths))])
+            
+            # Compute and store embeddings
+            embeddings = []
+            ids = []
+            metadatas = []
+            for i, img in enumerate(st.session_state.demo_images):
+                img_pre = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
+                with torch.no_grad():
+                    embedding = st.session_state.model.encode_image(img_pre).cpu().numpy().flatten()
+                embeddings.append(embedding)
+                ids.append(str(i))
+                metadatas.append({"path": demo_image_paths[i]})
+            
+            # Add to ChromaDB
+            try:
+                st.session_state.demo_collection.add(
+                    embeddings=embeddings,
+                    ids=ids,
+                    metadatas=metadatas
+                )
+            except Exception as e:
+                st.error(f"Failed to add demo images to ChromaDB: {e}")
         else:
             st.warning("No images found in 'demo_images' folder. Demo mode will be limited.")
 
@@ -45,26 +90,32 @@ if mode == "Search in My Images":
     uploaded_files = st.file_uploader("Choose images", type=['png', 'jpg', 'jpeg'], accept_multiple_files=True)
     
     if uploaded_files:
-        # Clear previous user images to avoid duplicates
+        # Clear_previous user images and collection
         st.session_state.user_images = []
-        st.session_state.user_encodings = []
+        st.session_state.user_collection.delete(ids=[str(i) for i in range(st.session_state.user_collection.count())])
         
-        for uploaded_file in uploaded_files:
+        for i, uploaded_file in enumerate(uploaded_files):
             img = Image.open(uploaded_file)
             st.session_state.user_images.append(img)
             img_pre = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
             with torch.no_grad():
-                encoding = st.session_state.model.encode_image(img_pre)
-            st.session_state.user_encodings.append(encoding)
+                embedding = st.session_state.model.encode_image(img_pre).cpu().numpy().flatten()
+            try:
+                st.session_state.user_collection.add(
+                    embeddings=[embedding],
+                    ids=[str(i)],
+                    metadatas=[{"index": i}]
+                )
+            except Exception as e:
+                st.error(f"Failed to add user image {i} to ChromaDB: {e}")
         
-        if st.session_state.user_encodings:
-            st.session_state.user_encodings = torch.cat(st.session_state.user_encodings, dim=0)
+        if st.session_state.user_collection.count() > 0:
             st.success(f"Uploaded {len(st.session_state.user_images)} images successfully.")
         else:
             st.warning("No images uploaded yet.")
 
 # Query image upload
-st.subheader("Upload Query Image")
+st.subheader Snip: st.subheader("Upload Query Image")
 query_file = st.file_uploader("Choose a query image", type=['png', 'jpg', 'jpeg'])
 
 if query_file is not None:
@@ -72,30 +123,44 @@ if query_file is not None:
     st.image(query_img, caption="Query Image", width=200)
     query_pre = st.session_state.preprocess(query_img).unsqueeze(0).to(st.session_state.device)
     with torch.no_grad():
-        query_encoding = st.session_state.model.encode_image(query_pre)
+        query_embedding = st.session_state.model.encode_image(query_pre).cpu().numpy().flatten()
     
     if mode == "Search in Demo Images":
-        if st.session_state.demo_encodings is not None and len(st.session_state.demo_encodings) > 0:
-            similarities = (st.session_state.demo_encodings @ query_encoding.T).squeeze()
-            top_indices = torch.topk(similarities, min(5, len(similarities))).indices.cpu().numpy()
+        if st.session_state.demo_collection.count() > 0:
+            # Query ChromaDB
+            results = st.session_state.demo_collection.query(
+                query_embeddings=[query_embedding],
+                n_results=min(5, st.session_state.demo_collection.count())
+            )
+            distances = results['distances'][0]
+            ids = results['ids'][0]
+            similarities = [1 - dist for dist in distances]  # Convert distance to similarity
             
             st.subheader("Top 5 Similar Images")
             cols = st.columns(5)
-            for i, idx in enumerate(top_indices):
+            for i, (idx, sim) in enumerate(zip(ids, similarities)):
+                img_idx = int(idx)
                 with cols[i]:
-                    st.image(st.session_state.demo_images[idx], caption=f"Similarity: {similarities[idx]:.4f}", width=150)
+                    st.image(st.session_state.demo_images[img_idx], caption=f"Similarity: {sim:.4f}", width=150)
         else:
             st.error("No demo images available. Please check the 'demo_images' folder.")
     
     elif mode == "Search in My Images":
-        if st.session_state.user_encodings is not None and len(st.session_state.user_encodings) > 0:
-            similarities = (st.session_state.user_encodings @ query_encoding.T).squeeze()
-            top_indices = torch.topk(similarities, min(5, len(similarities))).indices.cpu().numpy()
+        if st.session_state.user_collection.count() > 0:
+            # Query ChromaDB
+            results = st.session_state.user_collection.query(
+                query_embeddings=[query_embedding],
+                n_results=min(5, st.session_state.user_collection.count())
+            )
+            distances = results['distances'][0]
+            ids = results['ids'][0]
+            similarities = [1 - dist for dist in distances]  # Convert distance to similarity
             
             st.subheader("Top 5 Similar Images")
             cols = st.columns(5)
-            for i, idx in enumerate(top_indices):
+            for i, (idx, sim) in enumerate(zip(ids, similarities)):
+                img_idx = int(idx)
                 with cols[i]:
-                    st.image(st.session_state.user_images[idx], caption=f"Similarity: {similarities[idx]:.4f}", width=150)
+                    st.image(st.session_state.user_images[img_idx], caption=f"Similarity: {sim:.4f}", width=150)
         else:
-            st.error("No user images uploaded yet. Please upload images first.")
+            st.error("No user images uploaded yet. Please upload images first.")