Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,26 +1,29 @@
 import os
-import fitz  # PyMuPDF
-import chromadb
+import fitz
 import tempfile
 import streamlit as st
 from PIL import Image
+from chromadb import PersistentClient
 from chromadb.utils.data_loaders import ImageLoader
 from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
+from skimage import data as skdata
+from skimage.io import imsave
+import uuid
 
 # Paths
 DB_PATH = './data/image_vdb'
 IMAGES_DIR = './data/extracted_images'
 os.makedirs(IMAGES_DIR, exist_ok=True)
 
-# Init Chroma
-chroma_client = chromadb.PersistentClient(path=DB_PATH)
+# Init ChromaDB
+chroma_client = PersistentClient(path=DB_PATH)
 image_loader = ImageLoader()
 embedding_fn = OpenCLIPEmbeddingFunction()
 image_collection = chroma_client.get_or_create_collection(
     name="image", embedding_function=embedding_fn, data_loader=image_loader
 )
 
-# Utilities
+# === Image Handling ===
 def extract_images_from_pdf(pdf_bytes):
     pdf = fitz.open(stream=pdf_bytes, filetype="pdf")
     saved_images = []
@@ -44,13 +47,12 @@ def extract_images_from_pdf(pdf_bytes):
 
     return saved_images
 
-def index_images_in_chroma(image_paths):
+def index_images(image_paths):
     ids = []
     uris = []
-
     for i, path in enumerate(sorted(image_paths)):
         if path.endswith((".png", ".jpeg", ".jpg")):
-            ids.append(f"img_{len(image_collection.get()['ids']) + i}")
+            ids.append(str(uuid.uuid4()))
             uris.append(path)
 
     if ids:
@@ -65,27 +67,74 @@ def query_similar_images(image_file, top_k=5):
     os.remove(tmp_path)
     return results['uris'][0]
 
-# Streamlit UI
-st.title("🔍 Image Search from PDF (HR Tool Demo)")
+def load_skimage_demo_images():
+    demo_images = {
+        "astronaut": skdata.astronaut(),
+        "coffee": skdata.coffee(),
+        "camera": skdata.camera(),
+        "chelsea": skdata.chelsea(),
+        "rocket": skdata.rocket()
+    }
+    saved_paths = []
+
+    for name, img in demo_images.items():
+        path = os.path.join(IMAGES_DIR, f"{name}.png")
+        imsave(path, img)
+        saved_paths.append(path)
+
+    return saved_paths
+
+# === Streamlit UI ===
+st.title("🔍 Image Similarity Search from PDF or Custom Dataset")
+
+# Source Selector
+source = st.radio(
+    "Select Image Source",
+    ["Upload PDF", "Upload Images", "Load Demo Dataset"],
+    horizontal=True
+)
 
-with st.expander("📤 Step 1: Upload PDF to Extract Images"):
-    uploaded_pdf = st.file_uploader("Upload a PDF file", type=["pdf"])
-    if uploaded_pdf is not None:
+if source == "Upload PDF":
+    uploaded_pdf = st.file_uploader("📤 Upload PDF", type=["pdf"])
+    if uploaded_pdf:
         with st.spinner("Extracting images..."):
-            saved_images = extract_images_from_pdf(uploaded_pdf.read())
-            index_images_in_chroma(saved_images)
-            st.success(f"Extracted and indexed {len(saved_images)} images.")
-            st.image(saved_images, caption="Extracted images", width=150)
-
+            images = extract_images_from_pdf(uploaded_pdf.read())
+            index_images(images)
+            st.success(f"{len(images)} images extracted and indexed.")
+            st.image(images, width=150)
+
+elif source == "Upload Images":
+    uploaded_imgs = st.file_uploader("📤 Upload one or more images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
+    if uploaded_imgs:
+        saved_paths = []
+        for img in uploaded_imgs:
+            img_path = os.path.join(IMAGES_DIR, img.name)
+            with open(img_path, "wb") as f:
+                f.write(img.read())
+            saved_paths.append(img_path)
+
+        index_images(saved_paths)
+        st.success(f"{len(saved_paths)} images indexed.")
+        st.image(saved_paths, width=150)
+
+elif source == "Load Demo Dataset":
+    if st.button("🔄 Load Demo Images (skimage)"):
+        demo_paths = load_skimage_demo_images()
+        index_images(demo_paths)
+        st.success("Demo images loaded and indexed.")
+        st.image(demo_paths, width=150)
+
+# Divider
 st.divider()
 
-with st.expander("🖼️ Step 2: Search by Uploading a Query Image"):
-    query_img = st.file_uploader("Upload a query image", type=["jpg", "jpeg", "png"])
-    if query_img is not None:
-        st.image(query_img, caption="Query Image", width=200)
-        with st.spinner("Searching similar images..."):
-            results = query_similar_images(query_img, top_k=5)
-
-        st.subheader("🔎 Top Matches:")
-        for res_path in results:
-            st.image(res_path, width=200, caption=os.path.basename(res_path))
+# Query Interface
+st.subheader("🔎 Search for Similar Images")
+query_img = st.file_uploader("Upload a query image", type=["jpg", "jpeg", "png"])
+if query_img:
+    st.image(query_img, caption="Query Image", width=200)
+    with st.spinner("Searching..."):
+        matches = query_similar_images(query_img, top_k=5)
+
+    st.subheader("📊 Top Matches:")
+    for match in matches:
+        st.image(match, width=200, caption=os.path.basename(match))