Building Image Search Engine with CLIP Embeddings + FAISS

# Traditional approach - manual tags
image_metadata = {
    "filename": "IMG_1234.jpg",
    "tags": ["mountain", "sunset", "landscape"],
    "date": "2024-08-15"
}
 
# Search query: "orange sky over snowy peaks"
# ❌ Won't find the image because exact tags don't match

import torch
import clip
from PIL import Image
import numpy as np
from tqdm import tqdm
import os
 
# Load CLIP model
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)
 
def embed_image(image_path: str) -> np.ndarray:
    """
    Generate 512-dimensional embedding for an image.
    """
    image = Image.open(image_path).convert("RGB")
    image_input = preprocess(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        image_features = model.encode_image(image_input)
        # Normalize to unit length for cosine similarity
        image_features /= image_features.norm(dim=-1, keepdim=True)
    
    return image_features.cpu().numpy()[0]
 
def index_images(image_dir: str, batch_size: int = 32):
    """
    Index all images in directory, processing in batches for efficiency.
    """
    image_paths = []
    for root, _, files in os.walk(image_dir):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
                image_paths.append(os.path.join(root, file))
    
    print(f"Found {len(image_paths)} images to index")
    
    embeddings = []
    image_ids = []
    
    for i in tqdm(range(0, len(image_paths), batch_size)):
        batch_paths = image_paths[i:i + batch_size]
        
        # Process batch
        images = [preprocess(Image.open(path).convert("RGB")) for path in batch_paths]
        image_batch = torch.stack(images).to(device)
        
        with torch.no_grad():
            features = model.encode_image(image_batch)
            features /= features.norm(dim=-1, keepdim=True)
        
        embeddings.append(features.cpu().numpy())
        image_ids.extend(batch_paths)
    
    # Concatenate all batches
    embeddings = np.vstack(embeddings)
    
    print(f"✅ Generated {embeddings.shape[0]} embeddings of dimension {embeddings.shape[1]}")
    
    return embeddings, image_ids
 
# Example: Index 100k images
embeddings, image_ids = index_images("/data/images", batch_size=32)

import faiss
import pickle
 
def build_faiss_index(embeddings: np.ndarray, nlist: int = 256):
    """
    Build FAISS index with IVF (Inverted File) for fast approximate search.
    
    Args:
        embeddings: (N, 512) array of image embeddings
        nlist: Number of clusters (√N is a good heuristic)
    """
    dimension = embeddings.shape[1]  # 512 for CLIP ViT-B/32
    
    # Quantizer for IVF
    quantizer = faiss.IndexFlatIP(dimension)  # Inner product (cosine similarity)
    
    # IVF index with nlist clusters
    index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT)
    
    # Train the index (learns cluster centroids)
    print("Training FAISS index...")
    index.train(embeddings)
    
    # Add all embeddings to index
    print("Adding embeddings to index...")
    index.add(embeddings)
    
    # Set nprobe (number of clusters to search)
    # Higher nprobe = more accurate but slower
    index.nprobe = 32  # Good balance between speed and accuracy
    
    print(f"✅ FAISS index built: {index.ntotal} vectors indexed")
    
    return index
 
# Build index
faiss_index = build_faiss_index(embeddings, nlist=256)
 
# Save index and image IDs
faiss.write_index(faiss_index, "image_search.index")
with open("image_ids.pkl", "wb") as f:
    pickle.dump(image_ids, f)
 
print("Index saved to disk")

def search_images(
    query: str,
    faiss_index: faiss.Index,
    image_ids: list,
    top_k: int = 50
) -> list:
    """
    Search for images matching the text query.
    """
    # Encode text query with CLIP
    text_input = clip.tokenize([query]).to(device)
    
    with torch.no_grad():
        text_features = model.encode_text(text_input)
        text_features /= text_features.norm(dim=-1, keepdim=True)
    
    query_embedding = text_features.cpu().numpy()
    
    # Search FAISS index
    start_time = time.time()
    similarities, indices = faiss_index.search(query_embedding, top_k)
    search_time = (time.time() - start_time) * 1000  # ms
    
    # Format results
    results = []
    for i, (idx, score) in enumerate(zip(indices[0], similarities[0])):
        results.append({
            "rank": i + 1,
            "image_id": image_ids[idx],
            "similarity": float(score),
            "search_time_ms": search_time
        })
    
    return results
 
# Example search
results = search_images(
    query="sunset over mountains",
    faiss_index=faiss_index,
    image_ids=image_ids,
    top_k=50
)
 
for result in results[:5]:
    print(f"Rank {result['rank']}: {result['image_id']} (score: {result['similarity']:.3f})")

def multimodal_search(
    text_query: str = None,
    image_query: str = None,
    text_weight: float = 0.5,
    image_weight: float = 0.5,
    top_k: int = 50
) -> list:
    """
    Search using text, image, or weighted combination of both.
    """
    embeddings = []
    
    # Text query
    if text_query:
        text_input = clip.tokenize([text_query]).to(device)
        with torch.no_grad():
            text_features = model.encode_text(text_input)
            text_features /= text_features.norm(dim=-1, keepdim=True)
        embeddings.append((text_features.cpu().numpy(), text_weight))
    
    # Image query (find similar images)
    if image_query:
        image = Image.open(image_query).convert("RGB")
        image_input = preprocess(image).unsqueeze(0).to(device)
        with torch.no_grad():
            image_features = model.encode_image(image_input)
            image_features /= image_features.norm(dim=-1, keepdim=True)
        embeddings.append((image_features.cpu().numpy(), image_weight))
    
    # Weighted average of embeddings
    query_embedding = sum(emb * weight for emb, weight in embeddings)
    query_embedding /= query_embedding.norm()  # Re-normalize
    
    # Search
    similarities, indices = faiss_index.search(query_embedding, top_k)
    
    results = []
    for idx, score in zip(indices[0], similarities[0]):
        results.append({
            "image_id": image_ids[idx],
            "similarity": float(score)
        })
    
    return results
 
# Example: Search for "red jacket" but show me images similar to this reference
results = multimodal_search(
    text_query="person wearing red jacket",
    image_query="/reference/red_jacket_example.jpg",
    text_weight=0.6,
    image_weight=0.4,
    top_k=50
)

import json
import boto3
import numpy as np
import faiss
import torch
import clip
import pickle
from typing import Dict, List
 
# Load CLIP model (cold start)
device = "cpu"  # Lambda doesn't have GPU
model, preprocess = clip.load("ViT-B/32", device=device)
 
# Load FAISS index from EFS
faiss_index = faiss.read_index("/mnt/efs/image_search.index")
with open("/mnt/efs/image_ids.pkl", "rb") as f:
    image_ids = pickle.load(f)
 
# S3 client for generating signed URLs
s3_client = boto3.client('s3')
BUCKET_NAME = "realroll-images"
 
def lambda_handler(event, context):
    """
    AWS Lambda handler for image search API.
    """
    try:
        # Parse request
        body = json.loads(event['body'])
        query = body.get('query', '')
        top_k = body.get('top_k', 50)
        
        if not query:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Query is required'})
            }
        
        # Encode text query
        text_input = clip.tokenize([query]).to(device)
        with torch.no_grad():
            text_features = model.encode_text(text_input)
            text_features /= text_features.norm(dim=-1, keepdim=True)
        
        query_embedding = text_features.cpu().numpy()
        
        # Search FAISS
        similarities, indices = faiss_index.search(query_embedding, top_k)
        
        # Generate signed URLs for images
        results = []
        for idx, score in zip(indices[0], similarities[0]):
            image_path = image_ids[idx]
            # Generate signed URL (valid for 1 hour)
            signed_url = s3_client.generate_presigned_url(
                'get_object',
                Params={'Bucket': BUCKET_NAME, 'Key': image_path},
                ExpiresIn=3600
            )
            
            results.append({
                "image_url": signed_url,
                "similarity": float(score),
                "image_id": image_path
            })
        
        return {
            'statusCode': 200,
            'headers': {
                'Content-Type': 'application/json',
                'Access-Control-Allow-Origin': '*'
            },
            'body': json.dumps({
                'success': True,
                'query': query,
                'count': len(results),
                'results': results
            })
        }
        
    except Exception as e:
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

// components/ImageSearch.tsx
'use client';
 
import { useState } from 'react';
import { Search, Loader2 } from 'lucide-react';
 
interface SearchResult {
  image_url: string;
  similarity: number;
  image_id: string;
}
 
export default function ImageSearch() {
  const [query, setQuery] = useState('');
  const [results, setResults] = useState<SearchResult[]>([]);
  const [loading, setLoading] = useState(false);
 
  const handleSearch = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!query.trim()) return;
    
    setLoading(true);
    
    try {
      const response = await fetch('https://api.realroll.com/search', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ query, top_k: 50 }),
      });
      
      const data = await response.json();
      
      if (data.success) {
        setResults(data.results);
      }
    } catch (error) {
      console.error('Search failed:', error);
    } finally {
      setLoading(false);
    }
  };
 
  return (
    <div className="max-w-7xl mx-auto p-6">
      {/* Search Bar */}
      <form onSubmit={handleSearch} className="mb-8">
        <div className="relative">
          <Search className="absolute left-4 top-1/2 -translate-y-1/2 text-gray-400 w-6 h-6" />
          <input
            type="text"
            value={query}
            onChange={(e) => setQuery(e.target.value)}
            placeholder='Try: "sunset over mountains" or "person wearing red jacket"'
            className="w-full pl-14 pr-4 py-5 text-lg border-2 rounded-xl focus:ring-2 focus:ring-blue-500"
          />
        </div>
        
        <div className="mt-4 flex gap-2">
          <button
            type="button"
            onClick={() => setQuery('sunset over mountains')}
            className="px-4 py-2 bg-gray-100 rounded-lg hover:bg-gray-200"
          >
            🌄 Sunset
          </button>
          <button
            type="button"
            onClick={() => setQuery('person wearing red jacket')}
            className="px-4 py-2 bg-gray-100 rounded-lg hover:bg-gray-200"
          >
            🧥 Red Jacket
          </button>
          <button
            type="button"
            onClick={() => setQuery('coffee on wooden table')}
            className="px-4 py-2 bg-gray-100 rounded-lg hover:bg-gray-200"
          >
            ☕ Coffee
          </button>
          <button
            type="button"
            onClick={() => setQuery('happy dog playing fetch')}
            className="px-4 py-2 bg-gray-100 rounded-lg hover:bg-gray-200"
          >
            🐕 Dog Playing
          </button>
        </div>
      </form>
 
      {/* Loading State */}
      {loading && (
        <div className="flex items-center justify-center py-12">
          <Loader2 className="w-10 h-10 animate-spin text-blue-500" />
          <span className="ml-3 text-gray-600">Searching 100k+ images...</span>
        </div>
      )}
 
      {/* Results */}
      {!loading && results.length > 0 && (
        <>
          <p className="text-gray-600 mb-4">
            Found {results.length} images matching "{query}"
          </p>
          
          <div className="grid grid-cols-2 md:grid-cols-4 lg:grid-cols-5 gap-4">
            {results.map((result, index) => (
              <div 
                key={index}
                className="group relative overflow-hidden rounded-lg border hover:shadow-xl transition-shadow"
              >
                <img
                  src={result.image_url}
                  alt={`Result ${index + 1}`}
                  className="w-full h-48 object-cover group-hover:scale-105 transition-transform"
                />
                
                {/* Similarity Score Overlay */}
                <div className="absolute bottom-0 left-0 right-0 bg-gradient-to-t from-black/70 to-transparent p-3">
                  <div className="flex items-center justify-between text-white text-sm">
                    <span>#{index + 1}</span>
                    <span className="font-semibold">
                      {(result.similarity * 100).toFixed(0)}% match
                    </span>
                  </div>
                </div>
              </div>
            ))}
          </div>
        </>
      )}
    </div>
  );
}

# Profiling search pipeline
import time
 
def profile_search(query: str):
    timings = {}
    
    # 1. Text encoding
    start = time.time()
    text_input = clip.tokenize([query]).to(device)
    with torch.no_grad():
        text_features = model.encode_text(text_input)
        text_features /= text_features.norm(dim=-1, keepdim=True)
    timings['clip_encoding'] = (time.time() - start) * 1000
    
    # 2. FAISS search
    start = time.time()
    query_embedding = text_features.cpu().numpy()
    similarities, indices = faiss_index.search(query_embedding, 50)
    timings['faiss_search'] = (time.time() - start) * 1000
    
    # 3. Result formatting
    start = time.time()
    results = format_results(indices, similarities)
    timings['formatting'] = (time.time() - start) * 1000
    
    timings['total'] = sum(timings.values())
    
    return timings
 
# Example
timings = profile_search("sunset over mountains")
print(f"CLIP Encoding:  {timings['clip_encoding']:.1f}ms")
print(f"FAISS Search:   {timings['faiss_search']:.1f}ms")
print(f"Formatting:     {timings['formatting']:.1f}ms")
print(f"Total Latency:  {timings['total']:.1f}ms")

def incremental_index_update(new_images: list):
    # Load existing index
    index = faiss.read_index("/mnt/efs/image_search.index")
    
    # Generate embeddings for new images
    new_embeddings = []
    for image_path in new_images:
        embedding = embed_image(image_path)
        new_embeddings.append(embedding)
    
    new_embeddings = np.vstack(new_embeddings)
    
    # Add to index
    index.add(new_embeddings)
    
    # Save updated index
    faiss.write_index(index, "/mnt/efs/image_search_new.index")
    
    # Atomic swap
    os.rename("/mnt/efs/image_search_new.index", "/mnt/efs/image_search.index")

# English query
search_images("sunset over mountains")
 
# Spanish query (same results!)
search_images("puesta de sol sobre montañas")
 
# French query
search_images("coucher de soleil sur les montagnes")
 
# All return similar top results because CLIP understands semantic meaning!

def reverse_image_search(query_image_path: str, top_k: int = 50):
    # Encode query image
    query_embedding = embed_image(query_image_path)
    
    # Search FAISS
    similarities, indices = faiss_index.search(
        query_embedding.reshape(1, -1),
        top_k
    )
    
    return format_results(indices, similarities)
 
# Example
similar_images = reverse_image_search("/uploads/user_image.jpg")

def negative_search(
    positive_query: str,
    negative_query: str,
    top_k: int = 50
):
    # Embed both queries
    pos_emb = embed_text(positive_query)
    neg_emb = embed_text(negative_query)
    
    # Subtract negative embedding
    query_embedding = pos_emb - 0.5 * neg_emb
    query_embedding /= np.linalg.norm(query_embedding)
    
    # Search
    similarities, indices = faiss_index.search(query_embedding, top_k)
    return format_results(indices, similarities)
 
# Example: Mountains but NOT snow
results = negative_search(
    positive_query="mountains",
    negative_query="snow",
    top_k=50
)

def compositional_search(queries: list, weights: list, top_k: int = 50):
    embeddings = [embed_text(q) for q in queries]
    
    # Weighted average
    query_embedding = sum(e * w for e, w in zip(embeddings, weights))
    query_embedding /= np.linalg.norm(query_embedding)
    
    similarities, indices = faiss_index.search(query_embedding, top_k)
    return format_results(indices, similarities)
 
# Example: 70% mountains + 30% lake
results = compositional_search(
    queries=["mountains", "lake"],
    weights=[0.7, 0.3],
    top_k=50
)