Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,6 +7,7 @@ import numpy as np
 import chromadb
 import requests
 import tempfile
+import time
 
 # ----- Setup -----
 CACHE_DIR = tempfile.gettempdir()
@@ -14,6 +15,16 @@ CHROMA_PATH = os.path.join(CACHE_DIR, "chroma_db")
 DEMO_DIR = os.path.join(CACHE_DIR, "demo_images")
 os.makedirs(DEMO_DIR, exist_ok=True)
 
+# ----- Initialize Session State -----
+if 'dataset_loaded' not in st.session_state:
+    st.session_state.dataset_loaded = False
+if 'dataset_name' not in st.session_state:
+    st.session_state.dataset_name = None
+if 'demo_images' not in st.session_state:
+    st.session_state.demo_images = []
+if 'user_images' not in st.session_state:
+    st.session_state.user_images = []
+
 # ----- Load CLIP Model -----
 if 'model' not in st.session_state:
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -32,33 +43,46 @@ if 'chroma_client' not in st.session_state:
         name="user_images", metadata={"hnsw:space": "cosine"}
     )
 
+# ----- Title -----
 st.title("🔍 CLIP-Based Image Search")
 
-# Dataset selection
+# ----- Dataset Buttons -----
 col1, col2 = st.columns(2)
-use_demo = col1.button("📦 Use Demo Images")
-upload_own = col2.button("📤 Upload Your Images")
-
-dataset_loaded = False
-dataset_name = None
-
-# ----- Handle Demo Images -----
-if use_demo:
+if col1.button("📦 Use Demo Images"):
+    st.session_state.dataset_name = "demo"
+    st.session_state.dataset_loaded = False
+
+if col2.button("📤 Upload Your Images"):
+    st.session_state.dataset_name = "user"
+    st.session_state.dataset_loaded = False
+
+# ----- Download + Embed Demo Images -----
+def download_image_with_retry(url, path, retries=3, delay=1.0):
+    for attempt in range(retries):
+        try:
+            r = requests.get(url, timeout=10)
+            if r.status_code == 200:
+                with open(path, 'wb') as f:
+                    f.write(r.content)
+                return True
+        except Exception as e:
+            time.sleep(delay)
+    return False
+
+if st.session_state.dataset_name == "demo" and not st.session_state.dataset_loaded:
     with st.spinner("Downloading and indexing demo images..."):
         st.session_state.demo_collection.delete(ids=[str(i) for i in range(50)])
-        demo_image_paths = []
-        demo_images = []
-
+        demo_image_paths, demo_images = [], []
         for i in range(50):
             path = os.path.join(DEMO_DIR, f"img_{i+1:02}.jpg")
             if not os.path.exists(path):
                 url = f"https://picsum.photos/seed/{i}/1024/768"
-                response = requests.get(url)
-                if response.status_code == 200:
-                    with open(path, "wb") as f:
-                        f.write(response.content)
-            demo_image_paths.append(path)
-            demo_images.append(Image.open(path).convert("RGB"))
+                download_image_with_retry(url, path)
+            try:
+                demo_images.append(Image.open(path).convert("RGB"))
+                demo_image_paths.append(path)
+            except:
+                continue  # skip corrupted
 
         embeddings, ids, metadatas = [], [], []
         for i, img in enumerate(demo_images):
@@ -71,13 +95,12 @@ if use_demo:
 
         st.session_state.demo_collection.add(embeddings=embeddings, ids=ids, metadatas=metadatas)
         st.session_state.demo_images = demo_images
-        dataset_loaded = True
-        dataset_name = "demo"
+        st.session_state.dataset_loaded = True
 
-    st.success("Demo images loaded!")
+    st.success("✅ Demo images loaded!")
 
-# ----- Handle User Uploads -----
-if upload_own:
+# ----- Upload User Images -----
+if st.session_state.dataset_name == "user" and not st.session_state.dataset_loaded:
     uploaded = st.file_uploader("Upload your images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
     if uploaded:
         st.session_state.user_collection.delete(ids=[
@@ -85,9 +108,11 @@ if upload_own:
         ])
         user_images = []
         for i, file in enumerate(uploaded):
-            img = Image.open(file).convert("RGB")
+            try:
+                img = Image.open(file).convert("RGB")
+            except:
+                continue
             user_images.append(img)
-
             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()
@@ -96,13 +121,12 @@ if upload_own:
             )
 
         st.session_state.user_images = user_images
-        st.success(f"{len(user_images)} images uploaded.")
-        dataset_loaded = True
-        dataset_name = "user"
+        st.session_state.dataset_loaded = True
+        st.success(f"✅ Uploaded {len(user_images)} images.")
 
-# ----- Search UI -----
-if dataset_loaded:
-    st.subheader("Search Section")
+# ----- Search Section -----
+if st.session_state.dataset_loaded:
+    st.subheader("🔎 Search Section")
     query_type = st.radio("Search by:", ("Text", "Image"))
 
     query_embedding = None
@@ -112,18 +136,19 @@ if dataset_loaded:
             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()
-    else:
-        img_file = st.file_uploader("Upload query image", type=["jpg", "jpeg", "png"])
-        if img_file:
-            query_img = Image.open(img_file).convert("RGB")
+
+    elif query_type == "Image":
+        query_file = st.file_uploader("Upload query image", type=["jpg", "jpeg", "png"], key="query_image")
+        if query_file:
+            query_img = Image.open(query_file).convert("RGB")
             st.image(query_img, caption="Query Image", width=200)
-            img_tensor = st.session_state.preprocess(query_img).unsqueeze(0).to(st.session_state.device)
+            query_tensor = 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(img_tensor).cpu().numpy().flatten()
+                query_embedding = st.session_state.model.encode_image(query_tensor).cpu().numpy().flatten()
 
     # ----- Perform Search -----
     if query_embedding is not None:
-        if dataset_name == "demo":
+        if st.session_state.dataset_name == "demo":
             collection = st.session_state.demo_collection
             images = st.session_state.demo_images
         else:
@@ -139,12 +164,12 @@ if dataset_loaded:
             distances = results["distances"][0]
             similarities = [1 - d for d in distances]
 
-            st.subheader("Top Matches")
+            st.subheader("🔗 Top Matches")
             cols = st.columns(len(ids))
             for i, (img_id, sim) in enumerate(zip(ids, similarities)):
                 with cols[i]:
                     st.image(images[int(img_id)], caption=f"Sim: {sim:.3f}", width=150)
         else:
-            st.warning("No images in the collection.")
+            st.warning("No indexed images to search.")
 else:
-    st.info("Please click on one of the options above to load a dataset.")
+    st.info("👆 Please select a dataset (Demo or Upload Images) to begin.")