Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,28 +1,25 @@
 import os
-import tempfile
-
-# Set cache directory to temp or app folder
-cache_dir = os.path.join(tempfile.gettempdir(), "hf_cache")
-os.makedirs(cache_dir, exist_ok=True)
-
-os.environ["XDG_CACHE_HOME"] = cache_dir
-os.environ["HF_HOME"] = cache_dir
-
-# Now import OpenCLIPEmbeddingFunction
-from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
-
+import uuid
 import fitz
 import tempfile
 import streamlit as st
 from PIL import Image
+import numpy as np
+from skimage.io import imsave
+from torchvision.datasets import CIFAR10
+import torchvision.transforms as T
+
+# Setup cache paths
+HF_CACHE = os.path.join(tempfile.gettempdir(), "hf_cache")
+os.makedirs(HF_CACHE, exist_ok=True)
+os.environ["XDG_CACHE_HOME"] = HF_CACHE
+os.environ["HF_HOME"] = HF_CACHE
+
 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
 
-# Use safe temp directories for Streamlit or restricted environments
+# Directories
 TEMP_DIR = tempfile.gettempdir()
 IMAGES_DIR = os.path.join(TEMP_DIR, "extracted_images")
 DB_PATH = os.path.join(TEMP_DIR, "image_vdb")
@@ -40,122 +37,101 @@ def get_chroma_collection():
 
 image_collection = get_chroma_collection()
 
-# === Image Extraction ===
+# — PDFs & Uploads —
 def extract_images_from_pdf(pdf_bytes):
     pdf = fitz.open(stream=pdf_bytes, filetype="pdf")
-    saved_images = []
-
-    for page_num in range(len(pdf)):
-        page = pdf.load_page(page_num)
-        images = page.get_images(full=True)
-
-        for img_idx, img in enumerate(images):
-            xref = img[0]
-            base_image = pdf.extract_image(xref)
-            img_bytes = base_image["image"]
-            ext = base_image["ext"]
-            filename = f"page_{page_num+1}_img_{img_idx+1}.{ext}"
-            path = os.path.join(IMAGES_DIR, filename)
-
-            with open(path, "wb") as f:
-                f.write(img_bytes)
-
-            saved_images.append(path)
-
-    return saved_images
-
-# === Indexing ===
-def index_images(image_paths):
-    ids = []
-    uris = []
-    for path in sorted(image_paths):
-        if path.lower().endswith((".png", ".jpeg", ".jpg")):
+    saved = []
+    for i in range(len(pdf)):
+        for img in pdf.load_page(i).get_images(full=True):
+            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"])
+            saved.append(path)
+    return saved
+
+def index_images(paths):
+    ids, uris = [], []
+    for path in sorted(paths):
+        if path.lower().endswith((".jpg",".jpeg",".png")):
             ids.append(str(uuid.uuid4()))
             uris.append(path)
-
     if ids:
         image_collection.add(ids=ids, uris=uris)
 
-# === Querying ===
+# — Queries —
 def query_similar_images(image_file, top_k=5):
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
         tmp.write(image_file.read())
-        tmp_path = tmp.name
-
-    try:
-        results = image_collection.query(query_uris=[tmp_path], n_results=top_k)
-        return results['uris'][0]
-    finally:
-        os.remove(tmp_path)
-
-# === Demo images ===
-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 = st.radio(
-    "Select Image Source",
-    ["Upload PDF", "Upload Images", "Load Demo Dataset"],
-    horizontal=True
-)
-
-if source == "Upload PDF":
-    uploaded_pdf = st.file_uploader("📤 Upload PDF", type=["pdf"])
-    if uploaded_pdf:
-        with st.spinner("Extracting images..."):
-            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)
+        tmp.flush()
+        res = image_collection.query(query_uris=[tmp.name], n_results=top_k)
+    os.remove(tmp.name)
+    return res['uris'][0]
 
-st.divider()
+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) —
+@st.cache_resource
+def load_demo_cifar10(n=500):
+    dataset = CIFAR10(root=TEMP_DIR, download=True, train=True)
+    transform = T.ToPILImage()
+    saved = []
+    for i in range(min(n, len(dataset))):
+        img, label = dataset[i]
+        if not isinstance(img, Image.Image):
+            img = transform(img)
+        path = os.path.join(IMAGES_DIR, f"cifar10_{i}_{label}.png")
+        img.save(path)
+        saved.append(path)
+    return saved
+
+# — UI Starts —
+st.title("🔍 Image & Text Similarity Search with 500‑Image Demo DB")
+
+choice = st.radio("Select data source", ["Upload PDF", "Upload Images", "Load CIFAR‑10 Demo"], horizontal=True)
+
+if choice=="Upload PDF":
+    pdf = st.file_uploader("📤 Upload PDF", 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")
+        st.image(imgs, width=120)
+
+elif choice=="Upload Images":
+    imgs = st.file_uploader("📤 Upload images", accept_multiple_files=True, type=["jpg","jpeg","png"])
+    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")
+        st.image(paths, width=120)
+
+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")
+        st.image(paths[:20], width=100)
 
-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)
+st.divider()
+st.subheader("🔎 Image-Based Search")
+q = st.file_uploader("Upload a query image", type=["jpg","jpeg","png"])
+if q:
+    st.image(q, caption="Query"); 
     with st.spinner("Searching..."):
-        matches = query_similar_images(query_img, top_k=5)
+        out = query_similar_images(q, top_k=5)
+    st.subheader("Top Image Matches")
+    for u in out: st.image(u, width=150)
 
-    st.subheader("📊 Top Matches:")
-    for match in matches:
-        st.image(match, width=200, caption=os.path.basename(match))
+st.divider()
+st.subheader("📝 Text-to-Image Semantic Search")
+txt = st.text_input("Enter description (e.g. 'a beach'):")
+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)