Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,9 +10,8 @@ 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"
-    # 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
+    os.makedirs(cache_dir, exist_ok=True)
     try:
         model, preprocess = clip.load("ViT-B/32", device=device, download_root=cache_dir)
     except Exception as e:
@@ -25,39 +24,31 @@ if 'model' not in st.session_state:
     st.session_state.demo_image_paths = []
     st.session_state.user_images = []
 
-# Initialize ChromaDB client
+# Initialize ChromaDB
 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
+            name="demo_images", metadata={"hnsw:space": "cosine"}
         )
         st.session_state.user_collection = st.session_state.chroma_client.get_or_create_collection(
-            name="user_images",
-            metadata={"hnsw:space": "cosine"}
+            name="user_images", metadata={"hnsw:space": "cosine"}
         )
     except Exception as e:
-        st.error(f"Failed to initialize ChromaDB collections: {e}")
+        st.error(f"Failed to initialize ChromaDB: {e}")
         st.stop()
 
-# Load demo images into ChromaDB
-if not st.session_state.demo_images:
+# Load demo images only once
+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.endswith(('.png', '.jpg', '.jpeg'))]
-        if len(demo_image_paths) > 0:
+        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) for path in demo_image_paths]
-            
-            # Clear existing demo collection to avoid duplicates
+            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))])
-            
-            # Compute and store embeddings
-            embeddings = []
-            ids = []
-            metadatas = []
+
+            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():
@@ -65,37 +56,38 @@ if not st.session_state.demo_images:
                 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
                 )
+                st.session_state.demo_images_loaded = True
             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.")
+            st.warning("No images found in 'demo_images' folder.")
+    else:
+        st.warning("Folder 'demo_images' does not exist.")
 
-# Streamlit UI
-st.title("Image Search with CLIP")
+# UI title
+st.title("🔍 Image Search with CLIP")
 
 # Mode selection
 mode = st.radio("Select mode", ("Search in Demo Images", "Search in My Images"))
 
-# User images upload
+# Upload user 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:
-        # Clear_previous user images and collection
         st.session_state.user_images = []
         st.session_state.user_collection.delete(ids=[str(i) for i in range(st.session_state.user_collection.count())])
-        
+
         for i, uploaded_file in enumerate(uploaded_files):
-            img = Image.open(uploaded_file)
+            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():
@@ -107,60 +99,58 @@ if mode == "Search in My Images":
                     metadatas=[{"index": i}]
                 )
             except Exception as e:
-                st.error(f"Failed to add user image {i} to ChromaDB: {e}")
-        
+                st.error(f"Failed to add image {i}: {e}")
+
         if st.session_state.user_collection.count() > 0:
-            st.success(f"Uploaded {len(st.session_state.user_images)} images successfully.")
+            st.success(f"Uploaded {len(st.session_state.user_images)} images.")
         else:
-            st.warning("No images uploaded yet.")
+            st.warning("Upload failed.")
 
-# Query image upload
-st.subheader Snip: st.subheader("Upload 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)
+    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()
-    
+
     if mode == "Search in Demo Images":
         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")
+            similarities = [1 - dist for dist in distances]
+
+            st.subheader("Top 5 Similar Demo Images")
             cols = st.columns(5)
             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)
         else:
-            st.error("No demo images available. Please check the 'demo_images' folder.")
-    
+            st.error("No demo images available.")
+
     elif mode == "Search in My Images":
         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")
+            similarities = [1 - dist for dist in distances]
+
+            st.subheader("Top 5 Similar Uploaded Images")
             cols = st.columns(5)
             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)
         else:
-            st.error("No user images uploaded yet. Please upload images first.")
+            st.error("Please upload some images first.")