text_search.py

# *** USAGE ***
# > uv run python text_search.py --query "a birthday party for primary school kids that conveys fun" --class-name "birthday_party"

#!/usr/bin/env python3
"""
Text-to-Image Semantic Search Script

This script uses CLIP to calculate cosine similarities between text queries
and pre-extracted image embeddings, enabling semantic image search.

Usage:
    python text_search.py --query "beautiful sunset over mountains"
    
    or with uv:
    uv run text_search.py --query "romantic winter holiday"
"""

import os
import torch
import json
import argparse
from pathlib import Path
from transformers import CLIPProcessor, CLIPModel
from tqdm import tqdm
import warnings
import numpy as np
from datetime import datetime

# Suppress warnings for cleaner output
warnings.filterwarnings("ignore")

# Force use of safetensors to avoid torch.load security issue
os.environ["SAFETENSORS_FAST_GPU"] = "1"


class TextImageSearcher:
    """Calculate semantic similarities between text queries and image embeddings"""
    
    def __init__(self, model_name="openai/clip-vit-base-patch32", device=None):
        """
        Initialize the text-image searcher
        
        Args:
            model_name (str): Hugging Face model name for CLIP
            device (str): Device to run model on ('cuda', 'mps', 'cpu', or None for auto)
        """
        self.model_name = model_name
        self.device = self._get_device(device)
        
        print(f"🔍 Loading CLIP model for text search: {model_name}")
        print(f"🔧 Using device: {self.device}")
        
        # Load CLIP model and processor (using safetensors to avoid torch.load security issue)
        self.model = CLIPModel.from_pretrained(model_name, use_safetensors=True)
        self.processor = CLIPProcessor.from_pretrained(model_name)
        
        # Move model to device
        self.model = self.model.to(self.device)
        self.model.eval()  # Set to evaluation mode
        
        print("✅ Model loaded successfully!")
    
    def _get_device(self, device=None):
        """Auto-detect best device if not specified"""
        if device is not None:
            return device
        
        if torch.cuda.is_available():
            return "cuda"
        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            return "mps"  # Apple Silicon
        else:
            return "cpu"
    
    def encode_text(self, text):
        """Encode text query using CLIP text encoder"""
        try:
            # Preprocess text
            inputs = self.processor(text=[text], return_tensors="pt", padding=True, truncation=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            # Extract text features
            with torch.no_grad():
                text_features = self.model.get_text_features(**inputs)
                # Normalize features (important for cosine similarity)
                text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True)
            
            return text_features.cpu()
        except Exception as e:
            print(f"❌ Error encoding text '{text}': {e}")
            return None
    
    def load_image_embeddings(self, embeddings_dir="./embeddings", class_name=None, model_type="clip"):
        """Load all image embeddings from the embeddings directory"""
        # Use hierarchical structure if class_name is provided
        if class_name:
            embeddings_path = Path(embeddings_dir) / class_name / model_type
        else:
            embeddings_path = Path(embeddings_dir)
        
        if not embeddings_path.exists():
            print(f"❌ Embeddings directory not found: {embeddings_dir}")
            return None, None
        
        # Load mapping file
        mapping_file = embeddings_path / "image_to_feature_mapping.json"
        if not mapping_file.exists():
            print(f"❌ Mapping file not found: {mapping_file}")
            return None, None
        
        with open(mapping_file, 'r') as f:
            image_to_embedding = json.load(f)
        
        print(f"📁 Loading {len(image_to_embedding)} image embeddings...")
        
        embeddings = []
        image_names = []
        failed_loads = 0
        
        for image_name, embedding_file in tqdm(image_to_embedding.items(), desc="Loading embeddings"):
            embedding_path = embeddings_path / embedding_file
            try:
                # Load embedding tensor
                embedding = torch.load(embedding_path, map_location='cpu')
                embeddings.append(embedding)
                image_names.append(image_name)
            except Exception as e:
                print(f"⚠️  Failed to load {embedding_file}: {e}")
                failed_loads += 1
        
        if not embeddings:
            print("❌ No embeddings could be loaded!")
            return None, None
        
        # Stack all embeddings into a single tensor
        embeddings_tensor = torch.cat(embeddings, dim=0)
        
        print(f"✅ Loaded {len(embeddings)} embeddings successfully")
        if failed_loads > 0:
            print(f"⚠️  Failed to load {failed_loads} embeddings")
        
        return embeddings_tensor, image_names
    
    def calculate_similarities(self, text_query, embeddings_dir="./embeddings", class_name=None, model_type="clip"):
        """
        Calculate cosine similarities between text query and all image embeddings
        
        Args:
            text_query (str): Text query to search for
            embeddings_dir (str): Directory containing image embeddings
            class_name (str): Class name for hierarchical structure
            model_type (str): Model type for hierarchical structure
            
        Returns:
            tuple: (similarities, image_names) or (None, None) if failed
        """
        # Encode text query
        text_embedding = self.encode_text(text_query)
        if text_embedding is None:
            return None, None
        
        # Load image embeddings
        image_embeddings, image_names = self.load_image_embeddings(embeddings_dir, class_name, model_type)
        if image_embeddings is None:
            return None, None
        
        print(f"🔍 Calculating similarities for query: '{text_query}'")
        
        # Calculate cosine similarities
        with torch.no_grad():
            # Cosine similarity = dot product of normalized vectors
            similarities = torch.mm(text_embedding, image_embeddings.t())
            similarities = similarities.squeeze().cpu().numpy()
        
        return similarities, image_names
    
    def get_sorted_indices(self, similarities):
        """
        Get indices that sort similarities from highest to lowest (most to least relevant)
        
        Args:
            similarities (numpy.ndarray): Similarity scores
            
        Returns:
            numpy.ndarray: Indices that sort similarities in descending order
        """
        # argsort gives ascending order, so we reverse it with [::-1] for descending
        sorted_indices = np.argsort(similarities)[::-1]
        return sorted_indices
    
    def get_top_k_results(self, similarities, image_names, k=10):
        """
        Get top-k most similar results with their indices
        
        Args:
            similarities (numpy.ndarray): Similarity scores
            image_names (list): Corresponding image names
            k (int): Number of top results to return
            
        Returns:
            tuple: (top_k_indices, top_k_similarities, top_k_image_names)
        """
        sorted_indices = self.get_sorted_indices(similarities)
        top_k_indices = sorted_indices[:k]
        top_k_similarities = similarities[top_k_indices]
        top_k_image_names = [image_names[i] for i in top_k_indices]
        
        return top_k_indices, top_k_similarities, top_k_image_names
    
    def save_search_results(self, query, similarities, image_names, scores_dir="./scores", class_name=None):
        """
        Save search results (similarities/scores) to files
        
        Args:
            query (str): The search query
            similarities (numpy.ndarray): Similarity scores
            image_names (list): Corresponding image names
            scores_dir (str): Directory to save scores
            class_name (str): Class name for hierarchical structure
        """
        # Use hierarchical structure if class_name is provided
        if class_name:
            scores_path = Path("./config/scores") / class_name
        else:
            scores_path = Path(scores_dir)
        scores_path.mkdir(parents=True, exist_ok=True)
        
        # Generate filename with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        safe_query = "".join(c for c in query if c.isalnum() or c in (' ', '-', '_')).strip()
        safe_query = safe_query.replace(' ', '_')[:50]  # Limit length
        
        # Save detailed results as JSON
        results = []
        for img_name, similarity in zip(image_names, similarities):
            results.append({
                "image_name": img_name,
                "similarity_score": float(similarity),
                "rank": None  # Will be filled after sorting
            })
        
        # Sort by similarity (highest first)
        results.sort(key=lambda x: x['similarity_score'], reverse=True)
        
        # Add ranks
        for i, result in enumerate(results):
            result['rank'] = i + 1
        
        # Prepare metadata
        metadata = {
            "query": query,
            "timestamp": datetime.now().isoformat(),
            "model_name": self.model_name,
            "device_used": self.device,
            "total_images": len(image_names),
            "max_similarity": float(np.max(similarities)),
            "min_similarity": float(np.min(similarities)),
            "mean_similarity": float(np.mean(similarities)),
            "std_similarity": float(np.std(similarities))
        }
        
        # Save full results
        results_file = scores_path / f"search_results_{safe_query}_{timestamp}.json"
        full_data = {
            "metadata": metadata,
            "results": results
        }
        
        with open(results_file, 'w') as f:
            json.dump(full_data, f, indent=2)
        
        # Save compact scores file (just scores array)
        scores_file = scores_path / f"scores_{safe_query}_{timestamp}.npy"
        np.save(scores_file, similarities)
        
        # Save sorted indices (most to least relevant)
        sorted_indices = self.get_sorted_indices(similarities)
        indices_file = scores_path / f"sorted_indices_{safe_query}_{timestamp}.npy"
        np.save(indices_file, sorted_indices)
        
        # Save image names mapping
        names_file = scores_path / f"image_names_{safe_query}_{timestamp}.json"
        with open(names_file, 'w') as f:
            json.dump(image_names, f, indent=2)
        
        print(f"\n✅ Search results saved!")
        print(f"📊 Full results: {results_file}")
        print(f"🎯 Scores array: {scores_file}")
        print(f"📈 Sorted indices: {indices_file}")
        print(f"🏷️  Image names: {names_file}")
        
        # Display top results
        print(f"\n🔥 Top 10 most similar images:")
        print("-" * 60)
        for i, result in enumerate(results[:10]):
            print(f"{i+1:2d}. {result['image_name']:<40} (score: {result['similarity_score']:.4f})")
        
        return results_file, scores_file, indices_file, names_file
    
    def search(self, text_query, embeddings_dir="./embeddings", scores_dir="./scores", save_results=True, class_name=None, model_type="clip"):
        """
        Complete search pipeline: encode text, calculate similarities, and optionally save results
        
        Args:
            text_query (str): Text query to search for
            embeddings_dir (str): Directory containing image embeddings
            scores_dir (str): Directory to save scores
            save_results (bool): Whether to save results to files
            class_name (str): Class name for hierarchical structure
            model_type (str): Model type for hierarchical structure
            
        Returns:
            tuple: (similarities, image_names, sorted_indices, results_file) or None if failed
                  - similarities: numpy array of similarity scores
                  - image_names: list of image names
                  - sorted_indices: numpy array of indices sorted by similarity (high to low)
                  - results_file: path to saved results file (or None if save_results=False)
        """
        similarities, image_names = self.calculate_similarities(text_query, embeddings_dir, class_name, model_type)
        
        if similarities is None:
            print("❌ Search failed!")
            return None
        
        print(f"🎯 Search completed! Found similarities for {len(similarities)} images")
        
        # Get sorted indices for convenience
        sorted_indices = self.get_sorted_indices(similarities)
        
        if save_results:
            results_file, scores_file, indices_file, names_file = self.save_search_results(
                text_query, similarities, image_names, scores_dir, class_name
            )
            return similarities, image_names, sorted_indices, results_file
        else:
            return similarities, image_names, sorted_indices, None


def main():
    """Main function for command-line usage"""
    parser = argparse.ArgumentParser(description="Text-to-Image Semantic Search using CLIP")
    parser.add_argument("--query", "-q", type=str, required=True, 
                       help="Text query to search for (e.g., 'beautiful sunset over mountains')")
    parser.add_argument("--embeddings-dir", type=str, default="./embeddings",
                       help="Directory containing image embeddings (default: ./embeddings)")
    parser.add_argument("--scores-dir", type=str, default="./scores",
                       help="Directory to save search scores (default: ./scores, or ./config/scores/<class-name> if class-name is provided)")
    parser.add_argument("--class-name", type=str, default=None,
                       help="Class name for hierarchical structure (e.g., 'cosmetic', 'party')")
    parser.add_argument("--model-type", type=str, default="clip",
                       help="Model type for hierarchical structure (default: clip)")
    parser.add_argument("--model", type=str, default="openai/clip-vit-base-patch32",
                       help="CLIP model to use (default: openai/clip-vit-base-patch32)")
    parser.add_argument("--device", type=str, default=None,
                       help="Device to use (cuda/mps/cpu, default: auto-detect)")
    parser.add_argument("--no-save", action="store_true",
                       help="Don't save results to files")
    
    args = parser.parse_args()
    
    print("🎨 Text-to-Image Semantic Search with CLIP")
    print("=" * 50)
    
    # Initialize searcher
    searcher = TextImageSearcher(
        model_name=args.model,
        device=args.device
    )
    
    # Perform search
    results = searcher.search(
        text_query=args.query,
        embeddings_dir=args.embeddings_dir,
        scores_dir=args.scores_dir,
        save_results=not args.no_save,
        class_name=args.class_name,
        model_type=args.model_type
    )
    
    if results is None:
        print("💥 Search failed!")
        return 1
    
    similarities, image_names, sorted_indices, results_file = results
    
    print("\n🎉 Search completed successfully!")
    return 0


if __name__ == "__main__":
    exit(main())