src/streamlit_app.py

import streamlit as st
import torch
import clip
from PIL import Image
import os
import numpy as np
import chromadb
import requests
import tempfile
import time

# ----- Setup -----
st.set_page_config(page_title="CLIP Image Search", layout="wide")
CACHE_DIR = tempfile.gettempdir()
CHROMA_PATH = os.path.join(CACHE_DIR, "chroma_db")
DEMO_DIR = os.path.join(CACHE_DIR, "demo_images")
os.makedirs(DEMO_DIR, exist_ok=True)

# ----- Session State Init -----
if 'dataset_loaded' not in st.session_state:
    st.session_state.dataset_loaded = False
if 'dataset_name' not in st.session_state:
    st.session_state.dataset_name = None
if 'demo_images' not in st.session_state:
    st.session_state.demo_images = []
if 'user_images' not in st.session_state:
    st.session_state.user_images = []

# ----- Load CLIP Model -----
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, download_root=CACHE_DIR)
    st.session_state.model = model
    st.session_state.preprocess = preprocess
    st.session_state.device = device

# ----- Initialize ChromaDB -----
if 'chroma_client' not in st.session_state:
    st.session_state.chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
    st.session_state.demo_collection = st.session_state.chroma_client.get_or_create_collection(
        name="demo_images", metadata={"hnsw:space": "cosine"}
    )
    st.session_state.user_collection = st.session_state.chroma_client.get_or_create_collection(
        name="user_images", metadata={"hnsw:space": "cosine"}
    )

# ----- Sidebar -----
with st.sidebar:
    st.title("🧠 CLIP Search App")
    st.markdown("Choose a dataset to begin:")
    if st.button("📦 Load Demo Images"):
        st.session_state.dataset_name = "demo"
        st.session_state.dataset_loaded = False
    if st.button("📤 Upload Your Images"):
        st.session_state.dataset_name = "user"
        st.session_state.dataset_loaded = False

# ----- Helper -----
def download_image_with_retry(url, path, retries=3, delay=1.0):
    for attempt in range(retries):
        try:
            r = requests.get(url, timeout=10)
            if r.status_code == 200:
                with open(path, 'wb') as f:
                    f.write(r.content)
                return True
        except Exception:
            time.sleep(delay)
    return False

# ----- Main App -----
left, right = st.columns([2, 1])

with left:
    st.title("🔍 CLIP-Based Image Search")

    # ----- Load Demo -----
    if st.session_state.dataset_name == "demo" and not st.session_state.dataset_loaded:
        with st.spinner("Downloading and indexing demo images..."):
            st.session_state.demo_collection.delete(ids=[str(i) for i in range(50)])
            demo_image_paths, demo_images = [], []
            for i in range(50):
                path = os.path.join(DEMO_DIR, f"img_{i+1:02}.jpg")
                if not os.path.exists(path):
                    url = f"https://picsum.photos/seed/{i}/1024/768"
                    download_image_with_retry(url, path)
                try:
                    demo_images.append(Image.open(path).convert("RGB"))
                    demo_image_paths.append(path)
                except:
                    continue
            embeddings, ids, metadatas = [], [], []
            for i, img in enumerate(demo_images):
                img_tensor = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
                with torch.no_grad():
                    embedding = st.session_state.model.encode_image(img_tensor).cpu().numpy().flatten()
                embeddings.append(embedding)
                ids.append(str(i))
                metadatas.append({"path": demo_image_paths[i]})
            st.session_state.demo_collection.add(embeddings=embeddings, ids=ids, metadatas=metadatas)
            st.session_state.demo_images = demo_images
            st.session_state.dataset_loaded = True
        st.success("✅ Demo images loaded!")

    # ----- Upload User Images -----
    if st.session_state.dataset_name == "user" and not st.session_state.dataset_loaded:
        uploaded = st.file_uploader("Upload your images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
        if uploaded:
            st.session_state.user_collection.delete(ids=[
                str(i) for i in range(st.session_state.user_collection.count())
            ])
            user_images = []
            for i, file in enumerate(uploaded):
                try:
                    img = Image.open(file).convert("RGB")
                except:
                    continue
                user_images.append(img)
                img_tensor = st.session_state.preprocess(img).unsqueeze(0).to(st.session_state.device)
                with torch.no_grad():
                    embedding = st.session_state.model.encode_image(img_tensor).cpu().numpy().flatten()
                st.session_state.user_collection.add(
                    embeddings=[embedding], ids=[str(i)], metadatas=[{"index": i}]
                )
            st.session_state.user_images = user_images
            st.session_state.dataset_loaded = True
            st.success(f"✅ Uploaded {len(user_images)} images.")

    # ----- Search Section -----
    if st.session_state.dataset_loaded:
        st.subheader("🔎 Search")
        query_type = st.radio("Search by:", ("Text", "Image"))

        query_embedding = None
        if query_type == "Text":
            text_query = st.text_input("Enter your search prompt:")
            if text_query:
                tokens = clip.tokenize([text_query]).to(st.session_state.device)
                with torch.no_grad():
                    query_embedding = st.session_state.model.encode_text(tokens).cpu().numpy().flatten()
        elif query_type == "Image":
            query_file = st.file_uploader("Upload query image", type=["jpg", "jpeg", "png"], key="query_image")
            if query_file:
                query_img = Image.open(query_file).convert("RGB")
                st.image(query_img, caption="Query Image", width=200)
                query_tensor = st.session_state.preprocess(query_img).unsqueeze(0).to(st.session_state.device)
                with torch.no_grad():
                    query_embedding = st.session_state.model.encode_image(query_tensor).cpu().numpy().flatten()

        # ----- Perform Search -----
        if query_embedding is not None:
            if st.session_state.dataset_name == "demo":
                collection = st.session_state.demo_collection
                images = st.session_state.demo_images
            else:
                collection = st.session_state.user_collection
                images = st.session_state.user_images

            if collection.count() > 0:
                results = collection.query(
                    query_embeddings=[query_embedding],
                    n_results=min(5, collection.count())
                )
                ids = results["ids"][0]
                distances = results["distances"][0]
                similarities = [1 - d for d in distances]

                st.subheader("🎯 Top Matches")
                cols = st.columns(len(ids))
                for i, (img_id, sim) in enumerate(zip(ids, similarities)):
                    with cols[i]:
                        st.image(images[int(img_id)], caption=f"Similarity: {sim:.3f}", use_column_width=True)
            else:
                st.warning("⚠️ No images available for search.")
    else:
        st.info("👈 Choose a dataset from the sidebar to get started.")

# ----- Right Panel: Show Current Dataset Images -----
with right:
    st.subheader("🖼️ Dataset Preview")
    image_list = st.session_state.demo_images if st.session_state.dataset_name == "demo" else st.session_state.user_images
    if st.session_state.dataset_loaded and image_list:
        st.caption(f"Showing {len(image_list)} images")
        for i, img in enumerate(image_list[:20]):
            st.image(img, use_column_width=True)
    else:
        st.markdown("No images to preview yet.")