Caption + Search: Image Search Pipeline

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import json
import os

# --- Configuration ---
EMBED_MODEL = "BAAI/bge-base-en-v1.5"  # 768-dim, strong retrieval
CAPTION_FILE = "captions.jsonl"          # Pre-generated captions
INDEX_FILE = "image_search.index"

# --- Step 2: Embed all captions ---
embed_model = SentenceTransformer(EMBED_MODEL)

# Load pre-generated captions (from Step 1)
# Each line: {"path": "images/001.jpg", "caption": "a dog..."}
entries = []
with open(CAPTION_FILE) as f:
    for line in f:
        entries.append(json.loads(line))

captions = [e["caption"] for e in entries]
paths = [e["path"] for e in entries]

# Encode all captions in one batch (GPU-accelerated)
print(f"Embedding {len(captions)} captions...")
embeddings = embed_model.encode(
    captions,
    normalize_embeddings=True,  # Required for cosine similarity
    show_progress_bar=True,
    batch_size=256
)
embeddings = embeddings.astype("float32")

# --- Step 3: Build FAISS index ---
dimension = embeddings.shape[1]  # 768 for bge-base
index = faiss.IndexFlatIP(dimension)  # Inner product = cosine sim
index.add(embeddings)

# Persist to disk
faiss.write_index(index, INDEX_FILE)
print(f"Index built: {index.ntotal} vectors, {dimension} dimensions")

# --- Search function ---
def search_images(query: str, k: int = 5) -> list[dict]:
    """Semantic search over captioned images.

    Args:
        query: Natural language search query
        k: Number of results to return

    Returns:
        List of {path, caption, score} dicts, sorted by relevance
    """
    query_vec = embed_model.encode(
        [query], normalize_embeddings=True
    ).astype("float32")

    scores, indices = index.search(query_vec, k)

    results = []
    for score, idx in zip(scores[0], indices[0]):
        if idx < 0:
            continue  # FAISS returns -1 for missing results
        results.append({
            "path": paths[idx],
            "caption": captions[idx],
            "score": round(float(score), 4)
        })
    return results

# --- Example ---
results = search_images("dog playing in the rain")
for r in results:
    print(f"  {r['score']:.3f}  {r['path']}")
    print(f"          {r['caption']}")
    print()

from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch
import faiss
import numpy as np

# Load CLIP
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

def embed_image(image_path):
    """Embed image directly into CLIP's joint space."""
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        features = model.get_image_features(**inputs)
    vec = features.numpy().flatten()
    return vec / np.linalg.norm(vec)

def embed_text(text):
    """Embed text query into CLIP's joint space."""
    inputs = processor(text=text, return_tensors="pt", padding=True)
    with torch.no_grad():
        features = model.get_text_features(**inputs)
    vec = features.numpy().flatten()
    return vec / np.linalg.norm(vec)

# Index images (no captioning needed)
image_vecs = np.array([embed_image(p) for p in image_paths])
index = faiss.IndexFlatIP(image_vecs.shape[1])  # 512-dim
index.add(image_vecs.astype("float32"))

# Search
query_vec = embed_text("dog playing in the rain")
D, I = index.search(query_vec.reshape(1, -1).astype("float32"), k=5)
for score, idx in zip(D[0], I[0]):
    print(f"  {score:.3f}  {image_paths[idx]}")