Blame for image_search.py

import sys
import platform
import os
import torch
import pickle
import numpy as np
from PIL import Image
from tqdm import tqdm
import clip  # ✅ Required for custom model queries using OpenAI's CLIP
from pathlib import Path

# SQLite patch for Streamlit deployment (non-Windows)
if platform.system() != "Windows":
    try:
        import pysqlite3
        sys.modules["sqlite3"] = pysqlite3
    except ImportError:
        raise ImportError("pysqlite3-binary is not installed. Add it to requirements.txt for deployment.")

# ✅ HuggingFace CLIP model for ChromaDB
from transformers import CLIPModel, CLIPProcessor

# ✅ ChromaDB and embedding functions
import chromadb
from chromadb.utils.embedding_functions import EmbeddingFunction

# Optional dataset downloader
from download import download_flickr8k_dataset

TOP_K = 5 

# -------------------------
# HuggingFace CLIP Wrapper
# -------------------------
class CLIPEmbeddingFunction(EmbeddingFunction):
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(self.device).eval()
        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    def embed_image(self, image):
        if isinstance(image, str):
            image = Image.open(image).convert("RGB")
        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
        with torch.no_grad():
            image_features = self.model.get_image_features(**inputs)
        return image_features.cpu().numpy()[0].tolist()

    def embed_text(self, text):
        inputs = self.processor(text=text, return_tensors="pt", padding=True, truncation=True).to(self.device)
        with torch.no_grad():
            text_features = self.model.get_text_features(**inputs)
        return text_features.cpu().numpy()[0].tolist()

    def embed_combined(self, text, image):
        text_emb = torch.tensor(self.embed_text(text))
        image_emb = torch.tensor(self.embed_image(image))
        avg_emb = (text_emb + image_emb) / 2
        return avg_emb.numpy().tolist()


# -------------------------
# CLIP Setup
# -------------------------
clip_embedder = CLIPEmbeddingFunction()

device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)


# -------------------------
# Load Custom Model Files
# -------------------------
with open("nn_model.pkl", "rb") as f:
    custom_nn_model = pickle.load(f)
with open("image_paths.pkl", "rb") as f:
    custom_image_paths = pickle.load(f)
custom_image_embeddings = np.load("image_embeddings.npy")


# -------------------------
# ChromaDB Training
# -------------------------
def train_rag_on_images(image_folder):
    if not os.path.exists(image_folder):
        print(f"{image_folder} not found. Downloading dataset...")
        download_flickr8k_dataset()

    client = chromadb.PersistentClient(path="./chroma_db")
    collection = client.get_or_create_collection(name="image_search")

    image_paths = []
    for root, _, files in os.walk(image_folder):
        for file in files:
            if file.lower().endswith((".jpg", ".jpeg", ".png")):
                image_paths.append(os.path.join(root, file))

    image_paths = image_paths[:4000]
    print(f"Embedding {len(image_paths)} images...")

    for idx, path in enumerate(tqdm(image_paths)):
        embedding = clip_embedder.embed_image(path)
        collection.add(
            documents=[path],
            embeddings=[embedding],
            ids=[f"img_{idx}"]
        )

    print("Training complete. Embeddings stored in ChromaDB.")


# -------------------------
# Unified Custom Query
# -------------------------
def query_custom_model(prompt=None, image=None, k=4):
    if prompt is None and image is None:
        raise ValueError("You must provide at least a prompt or an image.")

    query_embed = None

    if prompt:
        tokenized = clip.tokenize(prompt).to(device)
        with torch.no_grad():
            text_embed = clip_model.encode_text(tokenized).cpu().numpy()
        query_embed = text_embed

    if image:
        tensor = clip_preprocess(image).unsqueeze(0).to(device)
        with torch.no_grad():
            image_embed = clip_model.encode_image(tensor).cpu().numpy()
        if query_embed is not None:
            query_embed = (query_embed + image_embed) / 2.0
        else:
            query_embed = image_embed

    dists, indices = custom_nn_model.kneighbors(query_embed, n_neighbors=k)
    # return [custom_image_paths[i] for i in indices[0]]
    return [Path(custom_image_paths[i]).as_posix() for i in indices[0]]


# -------------------------
# ChromaDB Query
# -------------------------
def query_similar_images(prompt=None, image=None, k=4, show_custom=True):
    client = chromadb.PersistentClient(path="./chroma_db")
    collection = client.get_collection("image_search")

    if prompt and image:
        query_vector = clip_embedder.embed_combined(prompt, image)
    elif image:
        query_vector = clip_embedder.embed_image(image)
    elif prompt:
        query_vector = clip_embedder.embed_text(prompt)
    else:
        raise ValueError("Please provide at least a prompt or an image.")

    results = collection.query(query_embeddings=[query_vector], n_results=k)
    clip_image_paths = results['documents'][0]

    if any(not os.path.exists(path) for path in clip_image_paths):
        print("Some image files are missing. Re-downloading dataset...")
        download_flickr8k_dataset()

    custom_results = query_custom_model(prompt, image, k=k) if show_custom else []
    return clip_image_paths, custom_results


# -------------------------
# Optional Fallback Loader
# -------------------------
def load_custom_model(model_dir=".", base_image_dir="flickr8k/Flicker8k_Dataset"):
    with open(f"{model_dir}/nn_model.pkl", "rb") as f:
        nn_model = pickle.load(f)
    with open(f"{model_dir}/image_paths.pkl", "rb") as f:
        image_paths = pickle.load(f)
    image_embeddings = np.load(f"{model_dir}/image_embeddings.npy")

    # 🔧 Normalize paths for Linux environments
    image_paths = [p.replace("\\", "/") for p in image_paths]
    fixed_paths = []
    for path in image_paths:
        if not os.path.exists(path):
            filename = os.path.basename(path)
            fixed_paths.append(os.path.join(base_image_dir, filename))
        else:
            fixed_paths.append(path)

    return nn_model, image_embeddings, fixed_paths

# Inside image_search.py

def update_custom_image_paths(folder_path):
    global custom_image_paths
    custom_image_paths = [os.path.join(folder_path, name) for name in os.listdir(folder_path) if name.lower().endswith((".jpg", ".jpeg", ".png"))]

def query_by_image(image_path, nn_model, image_paths):
    image = Image.open(image_path).convert("RGB")
    tensor = clip_preprocess(image).unsqueeze(0).to(device)
    with torch.no_grad():
        image_features = clip_model.encode_image(tensor).cpu().numpy()
    dists, indices = nn_model.kneighbors(image_features, n_neighbors=TOP_K)
    return [image_paths[i] for i in indices[0]]

def query_by_text(prompt, nn_model, image_paths):
    # Tokenize and encode the text
    with torch.no_grad():
        text_tokens = clip.tokenize([prompt]).to(device)
        text_features = clip_model.encode_text(text_tokens).cpu().numpy()

    # Get top matches
    dists, indices = nn_model.kneighbors(text_features, n_neighbors=TOP_K)
    return [image_paths[i] for i in indices[0]]