Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,12 +10,13 @@ import tempfile
 import time
 
 # ----- Setup -----
+st.set_page_config(page_title="CLIP Image Search", layout="wide")
 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)
 
-# ----- Initialize Session State -----
+# ----- Session State Init -----
 if 'dataset_loaded' not in st.session_state:
     st.session_state.dataset_loaded = False
 if 'dataset_name' not in st.session_state:
@@ -43,20 +44,18 @@ if 'chroma_client' not in st.session_state:
         name="user_images", metadata={"hnsw:space": "cosine"}
     )
 
-# ----- Title -----
-st.title("🔍 CLIP-Based Image Search")
-
-# ----- Dataset Buttons -----
-col1, col2 = st.columns(2)
-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 -----
+# ----- Sidebar -----
+with st.sidebar:
+    st.title("🧠 CLIP Search App")
+    st.markdown("Choose a dataset to begin:")
+    if st.button("📦 Load Demo Images"):
+        st.session_state.dataset_name = "demo"
+        st.session_state.dataset_loaded = False
+    if st.button("📤 Upload Your Images"):
+        st.session_state.dataset_name = "user"
+        st.session_state.dataset_loaded = False
+
+# ----- Helper -----
 def download_image_with_retry(url, path, retries=3, delay=1.0):
     for attempt in range(retries):
         try:
@@ -65,111 +64,124 @@ def download_image_with_retry(url, path, retries=3, delay=1.0):
                 with open(path, 'wb') as f:
                     f.write(r.content)
                 return True
-        except Exception as e:
+        except Exception:
             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 = [], []
-        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"
-                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):
-            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]})
-
-        st.session_state.demo_collection.add(embeddings=embeddings, ids=ids, metadatas=metadatas)
-        st.session_state.demo_images = demo_images
-        st.session_state.dataset_loaded = True
-
-    st.success("✅ Demo images loaded!")
-
-# ----- 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=[
-            str(i) for i in range(st.session_state.user_collection.count())
-        ])
-        user_images = []
-        for i, file in enumerate(uploaded):
-            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()
-            st.session_state.user_collection.add(
-                embeddings=[embedding], ids=[str(i)], metadatas=[{"index": i}]
-            )
-
-        st.session_state.user_images = user_images
-        st.session_state.dataset_loaded = True
-        st.success(f"✅ Uploaded {len(user_images)} images.")
-
-# ----- Search Section -----
-if st.session_state.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()
-
-    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)
-            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(query_tensor).cpu().numpy().flatten()
-
-    # ----- Perform Search -----
-    if query_embedding is not None:
-        if st.session_state.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 indexed images to search.")
-else:
-    st.info("👆 Please select a dataset (Demo or Upload Images) to begin.")
+# ----- Main App -----
+left, right = st.columns([2, 1])
+
+with left:
+    st.title("🔍 CLIP-Based Image Search")
+
+    # ----- Load Demo -----
+    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 = [], []
+            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"
+                    download_image_with_retry(url, path)
+                try:
+                    demo_images.append(Image.open(path).convert("RGB"))
+                    demo_image_paths.append(path)
+                except:
+                    continue
+            embeddings, ids, metadatas = [], [], []
+            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()
+                embeddings.append(embedding)
+                ids.append(str(i))
+                metadatas.append({"path": demo_image_paths[i]})
+            st.session_state.demo_collection.add(embeddings=embeddings, ids=ids, metadatas=metadatas)
+            st.session_state.demo_images = demo_images
+            st.session_state.dataset_loaded = True
+        st.success("✅ Demo images loaded!")
+
+    # ----- 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=[
+                str(i) for i in range(st.session_state.user_collection.count())
+            ])
+            user_images = []
+            for i, file in enumerate(uploaded):
+                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()
+                st.session_state.user_collection.add(
+                    embeddings=[embedding], ids=[str(i)], metadatas=[{"index": i}]
+                )
+            st.session_state.user_images = user_images
+            st.session_state.dataset_loaded = True
+            st.success(f"✅ Uploaded {len(user_images)} images.")
+
+    # ----- Search Section -----
+    if st.session_state.dataset_loaded:
+        st.subheader("🔎 Search")
+        query_type = st.radio("Search by:", ("Text", "Image"))
+
+        query_embedding = None
+        if query_type == "Text":
+            text_query = st.text_input("Enter your search prompt:")
+            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()
+        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)
+                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(query_tensor).cpu().numpy().flatten()
+
+        # ----- Perform Search -----
+        if query_embedding is not None:
+            if st.session_state.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"Similarity: {sim:.3f}", use_column_width=True)
+            else:
+                st.warning("⚠️ No images available for search.")
+    else:
+        st.info("👈 Choose a dataset from the sidebar to get started.")
+
+# ----- Right Panel: Show Current Dataset Images -----
+with right:
+    st.subheader("🖼️ Dataset Preview")
+    image_list = st.session_state.demo_images if st.session_state.dataset_name == "demo" else st.session_state.user_images
+    if st.session_state.dataset_loaded and image_list:
+        st.caption(f"Showing {len(image_list)} images")
+        for i, img in enumerate(image_list[:20]):
+            st.image(img, use_column_width=True)
+    else:
+        st.markdown("No images to preview yet.")