Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,20 +5,16 @@ from PIL import Image
 import os
 import numpy as np
 import chromadb
-from chromadb.utils import embedding_functions
 import tempfile
 
-# ----- Session Initialization -----
+# ----- Setup -----
+CACHE_DIR = tempfile.gettempdir()
+CHROMA_PATH = os.path.join(CACHE_DIR, "chroma_db")
+
+# ----- Load CLIP Model -----
 if 'model' not in st.session_state:
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    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()
-
+    model, preprocess = clip.load("ViT-B/32", device=device, download_root=CACHE_DIR)
     st.session_state.model = model
     st.session_state.preprocess = preprocess
     st.session_state.device = device
@@ -26,21 +22,15 @@ if 'model' not in st.session_state:
     st.session_state.demo_image_paths = []
     st.session_state.user_images = []
 
-# ----- Initialize ChromaDB in Temp Dir -----
+# ----- Initialize ChromaDB -----
 if 'chroma_client' not in st.session_state:
-    try:
-        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"}
-        )
-        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: {e}")
-        st.stop()
+    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"}
+    )
+    st.session_state.user_collection = st.session_state.chroma_client.get_or_create_collection(
+        name="user_images", metadata={"hnsw:space": "cosine"}
+    )
 
 # ----- Load Demo Images -----
 if not st.session_state.get("demo_images_loaded", False):
@@ -48,130 +38,93 @@ if not st.session_state.get("demo_images_loaded", False):
     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'))]
+        st.session_state.demo_images = [Image.open(p).convert("RGB") for p in demo_image_paths]
+        st.session_state.demo_image_paths = demo_image_paths
 
-        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]
-
-            # 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):
-                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]})
-
-            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.")
-    else:
-        st.warning("Folder 'demo_images' does not exist.")
+        st.session_state.demo_collection.delete(ids=[str(i) for i in range(len(demo_image_paths))])
 
-# ----- UI -----
-st.title("🔍 Image Search with CLIP")
-mode = st.radio("Select mode", ("Search in Demo Images", "Search in My Images"))
+        embeddings, ids, metadatas = [], [], []
+        for i, img in enumerate(st.session_state.demo_images):
+            img_tensor = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
+            with torch.no_grad():
+                embedding = st.session_state.model.encode_image(img_tensor).cpu().numpy().flatten()
+            embeddings.append(embedding)
+            ids.append(str(i))
+            metadatas.append({"path": demo_image_paths[i]})
 
-# ----- 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)
+        st.session_state.demo_collection.add(embeddings=embeddings, ids=ids, metadatas=metadatas)
+        st.session_state.demo_images_loaded = True
 
-    if uploaded_files:
+# ----- UI -----
+st.title("🔎 CLIP Image Search (Text & Image)")
+mode = st.radio("Choose dataset to search in:", ("Demo Images", "My Uploaded Images"))
+query_type = st.radio("Query type:", ("Image", "Text"))
+
+# ----- Upload User Images -----
+if mode == "My Uploaded Images":
+    uploaded = st.file_uploader("Upload your images", type=['jpg', 'jpeg', 'png'], accept_multiple_files=True)
+    if uploaded:
         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")
+        for i, file in enumerate(uploaded):
+            img = Image.open(file).convert("RGB")
             st.session_state.user_images.append(img)
 
-            img_pre = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
+            img_tensor = 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],
-                    ids=[str(i)],
-                    metadatas=[{"index": i}]
-                )
-            except Exception as 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.")
-        else:
-            st.warning("Upload failed.")
-
-# ----- 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(
-                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]
-
-            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"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(
-                query_embeddings=[query_embedding],
-                n_results=min(5, st.session_state.user_collection.count())
+                embedding = st.session_state.model.encode_image(img_tensor).cpu().numpy().flatten()
+
+            st.session_state.user_collection.add(
+                embeddings=[embedding],
+                ids=[str(i)],
+                metadatas=[{"index": i}]
             )
-            distances = results['distances'][0]
-            ids = results['ids'][0]
-            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"Sim: {sim:.4f}", width=150)
-        else:
-            st.error("Please upload some images first.")
+
+        st.success(f"{len(uploaded)} images uploaded.")
+
+# ----- Perform Query -----
+query_embedding = None
+if query_type == "Image":
+    img_file = st.file_uploader("Upload query image", type=["jpg", "jpeg", "png"])
+    if img_file:
+        img = Image.open(img_file).convert("RGB")
+        st.image(img, caption="Query Image", width=200)
+        img_tensor = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
+        with torch.no_grad():
+            query_embedding = st.session_state.model.encode_image(img_tensor).cpu().numpy().flatten()
+elif query_type == "Text":
+    text_query = st.text_input("Enter search text:")
+    if text_query:
+        tokens = clip.tokenize([text_query]).to(st.session_state.device)
+        with torch.no_grad():
+            query_embedding = st.session_state.model.encode_text(tokens).cpu().numpy().flatten()
+
+# ----- Run Search -----
+if query_embedding is not None:
+    if mode == "Demo Images":
+        collection = st.session_state.demo_collection
+        images = st.session_state.demo_images
+    else:
+        collection = st.session_state.user_collection
+        images = st.session_state.user_images
+
+    if collection.count() > 0:
+        results = collection.query(
+            query_embeddings=[query_embedding],
+            n_results=min(5, collection.count())
+        )
+        ids = results["ids"][0]
+        distances = results["distances"][0]
+        similarities = [1 - d for d in distances]
+
+        st.subheader("Top Matches")
+        cols = st.columns(5)
+        for i, (img_id, sim) in enumerate(zip(ids, similarities)):
+            with cols[i]:
+                idx = int(img_id)
+                st.image(images[idx], caption=f"Sim: {sim:.3f}", width=150)
+    else:
+        st.warning("No images found in collection.")