Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,210 +1,103 @@
-import os
-import uuid
-import fitz
-import tempfile
 import streamlit as st
+import torch
+import clip
 from PIL import Image
+import os
 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
-from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
-
-# 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)
-os.makedirs(DB_PATH, exist_ok=True)
-
-# Init ChromaDB collection
-@st.cache_resource
-def get_chroma_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):
-    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):
-    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):
-    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 []
-
-# --- Text-to-Image search ---
-def search_images_by_text(text, top_k=5):
-    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 []
-
-# --- Load CIFAR-10 Demo Dataset (500 images) ---
-@st.cache_resource
-def load_demo_cifar10(n=500):
-    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")
-
-# 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 file", type=["pdf"])
-    if pdf:
-        with st.spinner("Extracting images from PDF..."):
-            imgs = extract_images_from_pdf(pdf.read())
-            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)
-    if imgs:
-        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":
-    if st.button("🔄 Load 500 CIFAR‑10 Images"):
-        with st.spinner("Loading CIFAR‑10 demo dataset..."):
-            paths = load_demo_cifar10(500)
-            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()
-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 Image", width=200)
-    with st.spinner("Finding similar images..."):
-        results = query_similar_images(q, top_k=5)
-    if not results:
-        st.warning("No similar images found.")
-    else:
-        st.subheader("🔁 Top Matches:")
-        for u in results:
-            st.image(u, width=150)
 
-st.divider()
-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("Finding images by semantic similarity..."):
-        results = search_images_by_text(txt, top_k=5)
-    if not results:
-        st.warning("No semantic matches found.")
-    else:
-        st.subheader("🔍 Semantic Matches:")
-        for u in results:
-            st.image(u, width=150)
+# Initialize session state
+if 'model' not in st.session_state:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, preprocess = clip.load("ViT-B/32", device=device)
+    st.session_state.model = model
+    st.session_state.preprocess = preprocess
+    st.session_state.device = device
+    st.session_state.demo_images = []
+    st.session_state.demo_encodings = []
+    st.session_state.demo_image_paths = []
+    st.session_state.user_images = []
+    st.session_state.user_encodings = []
+
+# Load demo images
+if not st.session_state.demo_images:
+    demo_folder = "demo_images"
+    if os.path.exists(demo_folder):
+        demo_image_paths = [os.path.join(demo_folder, f) for f in os.listdir(demo_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
+        if len(demo_image_paths) > 0:
+            st.session_state.demo_image_paths = demo_image_paths
+            st.session_state.demo_images = [Image.open(path) for path in demo_image_paths]
+            demo_preprocessed = [st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device) for img in st.session_state.demo_images]
+            with torch.no_grad():
+                demo_encodings = [st.session_state.model.encode_image(img) for img in demo_preprocessed]
+            st.session_state.demo_encodings = torch.cat(demo_encodings, dim=0)
+        else:
+            st.warning("No images found in 'demo_images' folder. Demo mode will be limited.")
+
+# Streamlit UI
+st.title("Image Search with CLIP")
+
+# Mode selection
+mode = st.radio("Select mode", ("Search in Demo Images", "Search in My Images"))
+
+# User images upload
+if mode == "Search in My Images":
+    st.subheader("Upload Your Images")
+    uploaded_files = st.file_uploader("Choose images", type=['png', 'jpg', 'jpeg'], accept_multiple_files=True)
+    
+    if uploaded_files:
+        # Clear previous user images to avoid duplicates
+        st.session_state.user_images = []
+        st.session_state.user_encodings = []
+        
+        for uploaded_file in uploaded_files:
+            img = Image.open(uploaded_file)
+            st.session_state.user_images.append(img)
+            img_pre = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
+            with torch.no_grad():
+                encoding = st.session_state.model.encode_image(img_pre)
+            st.session_state.user_encodings.append(encoding)
+        
+        if st.session_state.user_encodings:
+            st.session_state.user_encodings = torch.cat(st.session_state.user_encodings, dim=0)
+            st.success(f"Uploaded {len(st.session_state.user_images)} images successfully.")
+        else:
+            st.warning("No images uploaded yet.")
+
+# Query image upload
+st.subheader("Upload Query Image")
+query_file = st.file_uploader("Choose a query image", type=['png', 'jpg', 'jpeg'])
+
+if query_file is not None:
+    query_img = Image.open(query_file)
+    st.image(query_img, caption="Query Image", width=200)
+    query_pre = st.session_state.preprocess(query_img).unsqueeze(0).
+
+to(st.session_state.device)
+    with torch.no_grad():
+        query_encoding = st.session_state.model.encode_image(query_pre)
+    
+    if mode == "Search in Demo Images":
+        if st.session_state.demo_encodings is not None and len(st.session_state.demo_encodings) > 0:
+            similarities = (st.session_state.demo_encodings @ query_encoding.T).squeeze()
+            top_indices = torch.topk(similarities, min(5, len(similarities))).indices.cpu().numpy()
+            
+            st.subheader("Top 5 Similar Images")
+            cols = st.columns(5)
+            for i, idx in enumerate(top_indices):
+                with cols[i]:
+                    st.image(st.session_state.demo_images[idx], caption=f"Similarity: {similarities[idx]:.4f}", width=150)
+        else:
+            st.error("No demo images available. Please check the 'demo_images' folder.")
+    
+    elif mode == "Search in My Images":
+        if st.session_state.user_encodings is not None and len(st.session_state.user_encodings) > 0:
+            similarities = (st.session_state.user_encodings @ query_encoding.T).squeeze()
+            top_indices = torch.topk(similarities, min(5, len(similarities))).indices.cpu().numpy()
+            
+            st.subheader("Top 5 Similar Images")
+            cols = st.columns(5)
+            for i, idx in enumerate(top_indices):
+                with cols[i]:
+                    st.image(st.session_state.user_images[idx], caption=f"Similarity: {similarities[idx]:.4f}", width=150)
+        else:
+            st.error("No user images uploaded yet. Please upload images first.")