Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,17 +6,19 @@ import os
 import numpy as np
 import chromadb
 from chromadb.utils import embedding_functions
-import tempfile  # ✅ For safe writable cache directory
+import tempfile
 
-# Initialize session state
+# ----- Session Initialization -----
 if 'model' not in st.session_state:
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    cache_dir = tempfile.gettempdir()  # ✅ Use temp dir instead of restricted local path
+    cache_dir = tempfile.gettempdir()
+
     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
@@ -24,10 +26,12 @@ if 'model' not in st.session_state:
     st.session_state.demo_image_paths = []
     st.session_state.user_images = []
 
-# Initialize ChromaDB
+# ----- Initialize ChromaDB in Temp Dir -----
 if 'chroma_client' not in st.session_state:
     try:
-        st.session_state.chroma_client = chromadb.PersistentClient(path="./chroma_db")
+        chroma_path = os.path.join(tempfile.gettempdir(), "chroma_db")
+        st.session_state.chroma_client = chromadb.PersistentClient(path=chroma_path)
+
         st.session_state.demo_collection = st.session_state.chroma_client.get_or_create_collection(
             name="demo_images", metadata={"hnsw:space": "cosine"}
         )
@@ -38,15 +42,22 @@ if 'chroma_client' not in st.session_state:
         st.error(f"Failed to initialize ChromaDB: {e}")
         st.stop()
 
-# Load demo images once
+# ----- Load Demo Images -----
 if not st.session_state.get("demo_images_loaded", False):
     demo_folder = "demo_images"
     if os.path.exists(demo_folder):
-        demo_image_paths = [os.path.join(demo_folder, f) for f in os.listdir(demo_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+        demo_image_paths = [os.path.join(demo_folder, f) for f in os.listdir(demo_folder)
+                            if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+
         if demo_image_paths:
             st.session_state.demo_image_paths = demo_image_paths
             st.session_state.demo_images = [Image.open(path).convert("RGB") for path in demo_image_paths]
-            st.session_state.demo_collection.delete(ids=[str(i) for i in range(len(demo_image_paths))])
+
+            # Clear previous collection
+            try:
+                st.session_state.demo_collection.delete(ids=[str(i) for i in range(len(demo_image_paths))])
+            except:
+                pass  # Collection might be empty
 
             embeddings, ids, metadatas = [], [], []
             for i, img in enumerate(st.session_state.demo_images):
@@ -71,27 +82,34 @@ if not st.session_state.get("demo_images_loaded", False):
     else:
         st.warning("Folder 'demo_images' does not exist.")
 
-# UI title
+# ----- UI -----
 st.title("🔍 Image Search with CLIP")
-
-# Mode selection
 mode = st.radio("Select mode", ("Search in Demo Images", "Search in My Images"))
 
-# Upload user images
+# ----- Upload My Images -----
 if mode == "Search in My Images":
     st.subheader("Upload Your Images")
     uploaded_files = st.file_uploader("Choose images", type=['png', 'jpg', 'jpeg'], accept_multiple_files=True)
 
     if uploaded_files:
         st.session_state.user_images = []
-        st.session_state.user_collection.delete(ids=[str(i) for i in range(st.session_state.user_collection.count())])
+
+        # Clear user collection
+        try:
+            st.session_state.user_collection.delete(ids=[
+                str(i) for i in range(st.session_state.user_collection.count())
+            ])
+        except:
+            pass
 
         for i, uploaded_file in enumerate(uploaded_files):
             img = Image.open(uploaded_file).convert("RGB")
             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():
                 embedding = st.session_state.model.encode_image(img_pre).cpu().numpy().flatten()
+
             try:
                 st.session_state.user_collection.add(
                     embeddings=[embedding],
@@ -106,17 +124,19 @@ if mode == "Search in My Images":
         else:
             st.warning("Upload failed.")
 
-# Query image
+# ----- Query Image -----
 st.subheader("Upload Query Image")
 query_file = st.file_uploader("Choose a query image", type=['png', 'jpg', 'jpeg'])
 
 if query_file is not None:
     query_img = Image.open(query_file).convert("RGB")
     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_embedding = st.session_state.model.encode_image(query_pre).cpu().numpy().flatten()
 
+    # ----- Search in Demo -----
     if mode == "Search in Demo Images":
         if st.session_state.demo_collection.count() > 0:
             results = st.session_state.demo_collection.query(
@@ -132,10 +152,11 @@ if query_file is not None:
             for i, (idx, sim) in enumerate(zip(ids, similarities)):
                 img_idx = int(idx)
                 with cols[i]:
-                    st.image(st.session_state.demo_images[img_idx], caption=f"Similarity: {sim:.4f}", width=150)
+                    st.image(st.session_state.demo_images[img_idx], caption=f"Sim: {sim:.4f}", width=150)
         else:
             st.error("No demo images available.")
 
+    # ----- Search in User Uploads -----
     elif mode == "Search in My Images":
         if st.session_state.user_collection.count() > 0:
             results = st.session_state.user_collection.query(
@@ -151,6 +172,6 @@ if query_file is not None:
             for i, (idx, sim) in enumerate(zip(ids, similarities)):
                 img_idx = int(idx)
                 with cols[i]:
-                    st.image(st.session_state.user_images[img_idx], caption=f"Similarity: {sim:.4f}", width=150)
+                    st.image(st.session_state.user_images[img_idx], caption=f"Sim: {sim:.4f}", width=150)
         else:
             st.error("Please upload some images first.")