Update src/streamlit_app.py

src/streamlit_app.py

@@ -9,7 +9,7 @@ from skimage.io import imsave
 from torchvision.datasets import CIFAR10
 import torchvision.transforms as T
 
-# Setup cache paths
+# Set HuggingFace cache directory
 HF_CACHE = os.path.join(tempfile.gettempdir(), "hf_cache")
 os.makedirs(HF_CACHE, exist_ok=True)
 os.environ["XDG_CACHE_HOME"] = HF_CACHE
@@ -19,12 +19,13 @@ from chromadb import PersistentClient
 from chromadb.utils.data_loaders import ImageLoader
 from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
 
-# Directories
+# Paths
 TEMP_DIR = tempfile.gettempdir()
 IMAGES_DIR = os.path.join(TEMP_DIR, "extracted_images")
 DB_PATH = os.path.join(TEMP_DIR, "image_vdb")
 os.makedirs(IMAGES_DIR, exist_ok=True)
 
+# Init ChromaDB collection
 @st.cache_resource
 def get_chroma_collection():
     chroma_client = PersistentClient(path=DB_PATH)
@@ -37,7 +38,7 @@ def get_chroma_collection():
 
 image_collection = get_chroma_collection()
 
-# — PDFs & Uploads —
+# --- Extract images from PDF ---
 def extract_images_from_pdf(pdf_bytes):
     pdf = fitz.open(stream=pdf_bytes, filetype="pdf")
     saved = []
@@ -46,20 +47,22 @@ def extract_images_from_pdf(pdf_bytes):
             base = pdf.extract_image(img[0])
             ext = base["ext"]
             path = os.path.join(IMAGES_DIR, f"pdf_p{i+1}_img{img[0]}.{ext}")
-            with open(path,"wb") as f: f.write(base["image"])
+            with open(path, "wb") as f:
+                f.write(base["image"])
             saved.append(path)
     return saved
 
+# --- Index images ---
 def index_images(paths):
     ids, uris = [], []
     for path in sorted(paths):
-        if path.lower().endswith((".jpg",".jpeg",".png")):
+        if path.lower().endswith((".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp")):
             ids.append(str(uuid.uuid4()))
             uris.append(path)
     if ids:
         image_collection.add(ids=ids, uris=uris)
 
-# — Queries —
+# --- Image-to-Image search ---
 def query_similar_images(image_file, top_k=5):
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
         tmp.write(image_file.read())
@@ -68,16 +71,19 @@ def query_similar_images(image_file, top_k=5):
     os.remove(tmp.name)
     return res['uris'][0]
 
+# --- Text-to-Image search ---
 def search_images_by_text(text, top_k=5):
     res = image_collection.query(query_texts=[text], n_results=top_k)
     return res['uris'][0]
 
-# — Demo Dataset: CIFAR10 (500 images) —
+# --- Load CIFAR-10 Demo Dataset (500 images) ---
 @st.cache_resource
 def load_demo_cifar10(n=500):
     dataset = CIFAR10(root=TEMP_DIR, download=True, train=True)
     transform = T.ToPILImage()
     saved = []
+
+    progress_bar = st.progress(0)
     for i in range(min(n, len(dataset))):
         img, label = dataset[i]
         if not isinstance(img, Image.Image):
@@ -85,53 +91,66 @@ def load_demo_cifar10(n=500):
         path = os.path.join(IMAGES_DIR, f"cifar10_{i}_{label}.png")
         img.save(path)
         saved.append(path)
+        if i % 10 == 0 or i == n - 1:
+            progress_bar.progress((i + 1) / n)
+
     return saved
 
-# — UI Starts —
-st.title("🔍 Image & Text Similarity Search with 500‑Image Demo DB")
+# === UI START ===
+st.title("🔍 Semantic Image Search App")
 
-choice = st.radio("Select data source", ["Upload PDF", "Upload Images", "Load CIFAR‑10 Demo"], horizontal=True)
+# Step 1: Load data
+choice = st.radio("📂 Select Image Source", ["Upload PDF", "Upload Images", "Load CIFAR‑10 Demo"], horizontal=True)
 
-if choice=="Upload PDF":
-    pdf = st.file_uploader("📤 Upload PDF", type=["pdf"])
+if choice == "Upload PDF":
+    pdf = st.file_uploader("📤 Upload PDF file", type=["pdf"])
     if pdf:
-        with st.spinner("Extracting..."):
-            imgs = extract_images_from_pdf(pdf.read()); index_images(imgs)
-        st.success(f"{len(imgs)} images indexed from PDF")
+        with st.spinner("Extracting images from PDF..."):
+            imgs = extract_images_from_pdf(pdf.read())
+            index_images(imgs)
+        st.success(f"✅ Indexed {len(imgs)} images from PDF.")
         st.image(imgs, width=120)
 
-elif choice=="Upload Images":
-    imgs = st.file_uploader("📤 Upload images", accept_multiple_files=True, type=["jpg","jpeg","png"])
+elif choice == "Upload Images":
+    imgs = st.file_uploader("📤 Upload image files", type=["jpg", "jpeg", "png", "bmp", "tiff", "webp"], accept_multiple_files=True)
     if imgs:
-        paths=[] 
-        for item in imgs:
-            p=os.path.join(IMAGES_DIR, item.name)
-            with open(p,"wb") as f: f.write(item.read()); paths.append(p)
-        index_images(paths)
-        st.success(f"{len(paths)} images uploaded & indexed")
+        with st.spinner("Indexing uploaded images..."):
+            paths = []
+            for item in imgs:
+                p = os.path.join(IMAGES_DIR, item.name)
+                with open(p, "wb") as f:
+                    f.write(item.read())
+                paths.append(p)
+            index_images(paths)
+        st.success(f"✅ {len(paths)} images indexed.")
         st.image(paths, width=120)
 
-elif choice=="Load CIFAR‑10 Demo":
+elif choice == "Load CIFAR‑10 Demo":
     if st.button("🔄 Load 500 CIFAR‑10 Images"):
-        paths=load_demo_cifar10(500); index_images(paths)
-        st.success("500 CIFAR‑10 demo images loaded and indexed")
+        with st.spinner("Loading CIFAR‑10 demo dataset..."):
+            paths = load_demo_cifar10(500)
+            index_images(paths)
+        st.success("✅ 500 demo images loaded and indexed.")
         st.image(paths[:20], width=100)
 
+# Step 2: Search
 st.divider()
-st.subheader("🔎 Image-Based Search")
-q = st.file_uploader("Upload a query image", type=["jpg","jpeg","png"])
+st.subheader("🖼️ Image-to-Image Search")
+q = st.file_uploader("📷 Upload a query image", type=["jpg", "jpeg", "png", "bmp", "tiff", "webp"])
 if q:
-    st.image(q, caption="Query"); 
-    with st.spinner("Searching..."):
-        out = query_similar_images(q, top_k=5)
-    st.subheader("Top Image Matches")
-    for u in out: st.image(u, width=150)
+    st.image(q, caption="Query Image", width=200)
+    with st.spinner("Finding similar images..."):
+        results = query_similar_images(q, top_k=5)
+    st.subheader("🔁 Top Matches:")
+    for u in results:
+        st.image(u, width=150)
 
 st.divider()
-st.subheader("📝 Text-to-Image Semantic Search")
-txt = st.text_input("Enter description (e.g. 'a beach'):")
+st.subheader("📝 Text-to-Image Search")
+txt = st.text_input("Describe what you’re looking for (e.g., 'a beach', 'a cat', 'a red truck'):")
 if txt:
-    with st.spinner("Searching..."):
-        out = search_images_by_text(txt, top_k=5)
-    st.subheader("Top Semantic Matches")
-    for u in out: st.image(u, width=150)
+    with st.spinner("Finding images by semantic similarity..."):
+        results = search_images_by_text(txt, top_k=5)
+    st.subheader("🔍 Semantic Matches:")
+    for u in results:
+        st.image(u, width=150)