Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,30 +7,12 @@ import numpy as np
 import chromadb
 import requests
 import tempfile
-from tqdm import tqdm
-
-# Get a temporary directory (automatically cleaned up after runtime ends)
-temp_dir = tempfile.gettempdir()
-demo_dir = os.path.join(temp_dir, "demo_images")
-os.makedirs(demo_dir, exist_ok=True)
-
-print(f"Saving images to: {demo_dir}")
-
-# Download 50 high-resolution images (1024x768)
-for i in tqdm(range(50), desc="Downloading images"):
-    url = f"https://picsum.photos/seed/{i}/1024/768"
-    response = requests.get(url)
-    if response.status_code == 200:
-        with open(os.path.join(demo_dir, f"img_{i+1:02}.jpg"), "wb") as f:
-            f.write(response.content)
-    else:
-        print(f"Failed to download image {i+1}")
-
-
 
 # ----- Setup -----
 CACHE_DIR = tempfile.gettempdir()
 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)
 
 # ----- Load CLIP Model -----
 if 'model' not in st.session_state:
@@ -39,9 +21,6 @@ if 'model' not in st.session_state:
     st.session_state.model = model
     st.session_state.preprocess = preprocess
     st.session_state.device = device
-    st.session_state.demo_images = []
-    st.session_state.demo_image_paths = []
-    st.session_state.user_images = []
 
 # ----- Initialize ChromaDB -----
 if 'chroma_client' not in st.session_state:
@@ -53,19 +32,36 @@ if 'chroma_client' not in st.session_state:
         name="user_images", metadata={"hnsw:space": "cosine"}
     )
 
-# ----- 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'))]
-        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
-
-        st.session_state.demo_collection.delete(ids=[str(i) for i in range(len(demo_image_paths))])
+st.title("🔍 CLIP-Based Image Search")
+
+# Dataset selection
+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:
+    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 = []
+
+        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"))
 
         embeddings, ids, metadatas = [], [], []
-        for i, img in enumerate(st.session_state.demo_images):
+        for i, img in enumerate(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()
@@ -74,78 +70,81 @@ if not st.session_state.get("demo_images_loaded", False):
             metadatas.append({"path": demo_image_paths[i]})
 
         st.session_state.demo_collection.add(embeddings=embeddings, ids=ids, metadatas=metadatas)
-        st.session_state.demo_images_loaded = True
+        st.session_state.demo_images = demo_images
+        dataset_loaded = True
+        dataset_name = "demo"
 
-# ----- 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"))
+    st.success("Demo images loaded!")
 
-# ----- Upload User Images -----
-if mode == "My Uploaded Images":
-    uploaded = st.file_uploader("Upload your images", type=['jpg', 'jpeg', 'png'], accept_multiple_files=True)
+# ----- Handle User Uploads -----
+if upload_own:
+    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())
         ])
-
+        user_images = []
         for i, file in enumerate(uploaded):
             img = Image.open(file).convert("RGB")
-            st.session_state.user_images.append(img)
+            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()
-
             st.session_state.user_collection.add(
-                embeddings=[embedding],
-                ids=[str(i)],
-                metadatas=[{"index": i}]
+                embeddings=[embedding], ids=[str(i)], metadatas=[{"index": i}]
             )
 
-        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)
+        st.session_state.user_images = user_images
+        st.success(f"{len(user_images)} images uploaded.")
+        dataset_loaded = True
+        dataset_name = "user"
+
+# ----- Search UI -----
+if dataset_loaded:
+    st.subheader("Search Section")
+    query_type = st.radio("Search by:", ("Text", "Image"))
+
+    query_embedding = None
+    if 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()
     else:
-        st.warning("No images found in collection.")
+        img_file = st.file_uploader("Upload query image", type=["jpg", "jpeg", "png"])
+        if img_file:
+            query_img = Image.open(img_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)
+            with torch.no_grad():
+                query_embedding = st.session_state.model.encode_image(img_tensor).cpu().numpy().flatten()
+
+    # ----- Perform Search -----
+    if query_embedding is not None:
+        if dataset_name == "demo":
+            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(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.")
+else:
+    st.info("Please click on one of the options above to load a dataset.")