Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,12 +8,19 @@ import numpy as np
 from skimage.io import imsave
 from torchvision.datasets import CIFAR10
 import torchvision.transforms as T
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
 
 # 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
 os.environ["HF_HOME"] = HF_CACHE
+# Add HuggingFace token if needed
+# os.environ["HF_TOKEN"] = "your-huggingface-api-token"
 
 from chromadb import PersistentClient
 from chromadb.utils.data_loaders import ImageLoader
@@ -24,83 +31,112 @@ 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)
+os.makedirs(DB_PATH, exist_ok=True)
 
 # Init ChromaDB collection
 @st.cache_resource
 def get_chroma_collection():
-    chroma_client = PersistentClient(path=DB_PATH)
-    image_loader = ImageLoader()
-    embedding_fn = OpenCLIPEmbeddingFunction()
-    collection = chroma_client.get_or_create_collection(
-        name="image", embedding_function=embedding_fn, data_loader=image_loader
-    )
-    return collection
+    try:
+        chroma_client = PersistentClient(path=DB_PATH)
+        image_loader = ImageLoader()
+        embedding_fn = OpenCLIPEmbeddingFunction()
+        collection = chroma_client.get_or_create_collection(
+            name="image", embedding_function=embedding_fn, data_loader=image_loader
+        )
+        return collection
+    except Exception as e:
+        logger.error(f"Error initializing ChromaDB: {e}")
+        st.error(f"Failed to initialize ChromaDB: {e}")
+        return None
 
 image_collection = get_chroma_collection()
+if image_collection is None:
+    st.stop()
 
 # --- Extract images from PDF ---
 def extract_images_from_pdf(pdf_bytes):
-    pdf = fitz.open(stream=pdf_bytes, filetype="pdf")
-    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
+    try:
+        pdf = fitz.open(stream=pdf_bytes, filetype="pdf")
+        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
+    except Exception as e:
+        logger.error(f"Error extracting images from PDF: {e}")
+        st.error(f"Failed to extract images: {e}")
+        return []
 
 # --- Index images ---
 def index_images(paths):
-    ids, uris = [], []
-    for path in sorted(paths):
-        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)
+    try:
+        ids, uris = [], []
+        for path in sorted(paths):
+            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)
+    except Exception as e:
+        logger.error(f"Error indexing images: {e}")
+        st.error(f"Failed to index images: {e}")
 
 # --- 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())
-        tmp.flush()
-        res = image_collection.query(query_uris=[tmp.name], n_results=top_k)
-    os.remove(tmp.name)
-    # Safe check for results
-    if not res or 'uris' not in res or not res['uris'] or not res['uris'][0]:
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
+            tmp.write(image_file.read())
+            tmp.flush()
+            res = image_collection.query(query_uris=[tmp.name], n_results=top_k)
+        os.remove(tmp.name)
+        if not res or 'uris' not in res or not res['uris'] or not res['uris'][0]:
+            return []
+        return res['uris'][0]
+    except Exception as e:
+        logger.error(f"Error in image-to-image search: {e}")
+        st.error(f"Failed to perform image search: {e}")
         return []
-    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)
-    # Safe check for results
-    if not res or 'uris' not in res or not res['uris'] or not res['uris'][0]:
+    try:
+        res = image_collection.query(query_texts=[text], n_results=top_k)
+        if not res or 'uris' not in res or not res['uris'] or not res['uris'][0]:
+            return []
+        return res['uris'][0]
+    except Exception as e:
+        logger.error(f"Error in text-to-image search: {e}")
+        st.error(f"Failed to perform text search: {e}")
         return []
-    return res['uris'][0]
 
 # --- 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):
-            img = transform(img)
-        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
+    try:
+        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):
+                img = transform(img)
+            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
+    except Exception as e:
+        logger.error(f"Error loading CIFAR-10 dataset: {e}")
+        st.error(f"Failed to load CIFAR-10 dataset: {e}")
+        return []
 
 # === UI START ===
 st.title("🔍 Semantic Image Search App")
@@ -113,9 +149,12 @@ if choice == "Upload PDF":
     if 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)
+            if imgs:
+                index_images(imgs)
+                st.success(f"✅ Indexed {len(imgs)} images from PDF.")
+                st.image(imgs, width=120)
+            else:
+                st.warning("No images extracted from PDF.")
 
 elif choice == "Upload Images":
     imgs = st.file_uploader("📤 Upload image files", type=["jpg", "jpeg", "png", "bmp", "tiff", "webp"], accept_multiple_files=True)
@@ -135,9 +174,12 @@ elif choice == "Load CIFAR‑10 Demo":
     if st.button("🔄 Load 500 CIFAR‑10 Images"):
         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)
+            if paths:
+                index_images(paths)
+                st.success("✅ 500 demo images loaded and indexed.")
+                st.image(paths[:20], width=100)
+            else:
+                st.warning("Failed to load CIFAR-10 images.")
 
 # Step 2: Search
 st.divider()
@@ -165,4 +207,4 @@ if txt:
     else:
         st.subheader("🔍 Semantic Matches:")
         for u in results:
-            st.image(u, width=150)
+            st.image(u, width=150)