Upload evaluation/fashion_search.py with huggingface_hub

evaluation/fashion_search.py

@@ -0,0 +1,365 @@
+#!/usr/bin/env python3
+"""
+Fashion search system using multi-modal embeddings.
+This file implements a fashion search engine that allows searching for clothing items
+using text queries. It uses embeddings from the main model to calculate cosine similarities
+and return the most relevant items. The system pre-computes embeddings for all items
+in the dataset for fast search.
+"""
+
+import torch
+import numpy as np
+import pandas as pd
+from PIL import Image
+import matplotlib.pyplot as plt
+from sklearn.metrics.pairwise import cosine_similarity
+from transformers import CLIPProcessor, CLIPModel as CLIPModel_transformers
+import warnings
+import os
+from typing import List, Tuple, Union, Optional
+import argparse
+
+# Import custom models
+from color_model import CLIPModel as ColorModel
+from hierarchy_model import Model as HierarchyModel, HierarchyExtractor
+from main_model import CustomDataset
+import config
+
+warnings.filterwarnings("ignore")
+
+class FashionSearchEngine:
+    """
+    Fashion search engine using multi-modal embeddings with category emphasis
+    """
+    
+    def __init__(self, top_k: int = 10, max_items: int = 10000):
+        """
+        Initialize the fashion search engine        
+        Args:
+            top_k: Number of top results to return
+            max_items: Maximum number of items to process (for faster initialization)
+            hierarchy_weight: Weight for hierarchy/category dimensions (default: 2.0)
+            color_weight: Weight for color dimensions (default: 1.0)
+        """
+        self.device = config.device
+        self.top_k = top_k
+        self.max_items = max_items
+        self.color_dim = config.color_emb_dim
+        self.hierarchy_dim = config.hierarchy_emb_dim
+        
+        # Load models
+        self._load_models()
+        
+        # Load dataset
+        self._load_dataset()
+        
+        # Pre-compute embeddings for all items
+        self._precompute_embeddings()
+        
+        print("✅ Fashion Search Engine ready!")
+    
+    def _load_models(self):
+        """Load all required models"""
+        print("📦 Loading models...")
+        
+        # Load color model
+        color_checkpoint = torch.load(config.color_model_path, map_location=self.device, weights_only=True)
+        self.color_model = ColorModel(embed_dim=self.color_dim).to(self.device)
+        self.color_model.load_state_dict(color_checkpoint)
+        self.color_model.eval()
+        
+        # Load hierarchy model
+        hierarchy_checkpoint = torch.load(config.hierarchy_model_path, map_location=self.device)
+        self.hierarchy_classes = hierarchy_checkpoint.get('hierarchy_classes', [])
+        self.hierarchy_model = HierarchyModel(
+            num_hierarchy_classes=len(self.hierarchy_classes),
+            embed_dim=self.hierarchy_dim
+        ).to(self.device)
+        self.hierarchy_model.load_state_dict(hierarchy_checkpoint['model_state'])
+        
+        # Set hierarchy extractor
+        hierarchy_extractor = HierarchyExtractor(self.hierarchy_classes, verbose=False)
+        self.hierarchy_model.set_hierarchy_extractor(hierarchy_extractor)
+        self.hierarchy_model.eval()
+        
+        # Load main CLIP model - Use the trained model directly
+        self.main_model = CLIPModel_transformers.from_pretrained('laion/CLIP-ViT-B-32-laion2B-s34B-b79K')
+        
+        # Load the trained weights
+        checkpoint = torch.load(config.main_model_path, map_location=self.device)
+        if 'model_state_dict' in checkpoint:
+            self.main_model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            # Fallback: try to load as state dict directly
+            self.main_model.load_state_dict(checkpoint)
+            print("✅ Loaded model weights directly")
+        
+        self.main_model.to(self.device)
+        self.main_model.eval()
+        
+        # Load CLIP processor
+        self.clip_processor = CLIPProcessor.from_pretrained('laion/CLIP-ViT-B-32-laion2B-s34B-b79K')
+        
+        print(f"✅ Models loaded - Colors: {self.color_dim}D, Hierarchy: {self.hierarchy_dim}D")
+    
+    def _load_dataset(self):
+        """Load the fashion dataset"""
+        print("📊 Loading dataset...")
+        
+        # Load dataset
+        self.df = pd.read_csv(config.local_dataset_path)
+        self.df_clean = self.df.dropna(subset=[config.column_local_image_path])
+        
+        # Create dataset object
+        self.dataset = CustomDataset(self.df_clean)
+        self.dataset.set_training_mode(False)  # No augmentation for search
+        
+        print(f"✅ {len(self.df_clean)} items loaded for search")
+    
+    def _precompute_embeddings(self):
+        """Pre-compute embeddings for all items in the dataset"""
+        print("🔄 Pre-computing embeddings...")
+        
+        # OPTIMIZATION: Sample a subset for faster initialization
+        print(f"⚠️ Dataset too large ({len(self.dataset)} items). Using stratified sampling of 10 items per color-category combination.")
+        
+        # Stratified sampling by color-category combinations
+        sampled_df = self.df_clean.groupby([config.color_column, config.hierarchy_column]).sample(n=20, replace=False)
+        
+        # Get the original indices of sampled items
+        sampled_indices = sampled_df.index.tolist()
+        
+        all_embeddings = []
+        all_texts = []
+        all_colors = []
+        all_hierarchies = []
+        all_images = []
+        all_urls = []
+        
+        # Process in batches for efficiency
+        batch_size = 32
+        
+        # Add progress bar
+        from tqdm import tqdm
+        total_batches = (len(sampled_indices) + batch_size - 1) // batch_size
+        
+        for i in tqdm(range(0, len(sampled_indices), batch_size), 
+                     desc="Computing embeddings", 
+                     total=total_batches):
+            batch_end = min(i + batch_size, len(sampled_indices))
+            batch_items = []
+            
+            for j in range(i, batch_end):
+                try:
+                    # Use the original dataset with the sampled index
+                    original_idx = sampled_indices[j]
+                    image, text, color, hierarchy = self.dataset[original_idx]
+                    batch_items.append((image, text, color, hierarchy))
+                    all_texts.append(text)
+                    all_colors.append(color)
+                    all_hierarchies.append(hierarchy)
+                    all_images.append(self.df_clean.iloc[original_idx][config.column_local_image_path])
+                    all_urls.append(self.df_clean.iloc[original_idx][config.column_url_image])
+                except Exception as e:
+                    print(f"⚠️ Skipping item {j}: {e}")
+                    continue
+            
+            if not batch_items:
+                continue
+            
+            # Process batch
+            images = torch.stack([item[0] for item in batch_items]).to(self.device)
+            texts = [item[1] for item in batch_items]
+            
+            with torch.no_grad():
+                # Get embeddings from main model (text embeddings only)
+                text_inputs = self.clip_processor(text=texts, padding=True, return_tensors="pt")
+                text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
+                
+                # Create dummy images for the model
+                dummy_images = torch.zeros(len(texts), 3, 224, 224).to(self.device)
+                
+                outputs = self.main_model(**text_inputs, pixel_values=dummy_images)
+                embeddings = outputs.text_embeds.cpu().numpy()
+                
+                all_embeddings.extend(embeddings)
+        
+        self.all_embeddings = np.array(all_embeddings)
+        self.all_texts = all_texts
+        self.all_colors = all_colors
+        self.all_hierarchies = all_hierarchies
+        self.all_images = all_images
+        self.all_urls = all_urls
+        
+        print(f"✅ Pre-computed embeddings for {len(self.all_embeddings)} items")
+    
+    def search_by_text(self, query_text: str, filter_category: str = None) -> List[dict]:
+        """
+        Search for clothing items using text query
+        
+        Args:
+            query_text: Text description to search for
+            
+        Returns:
+            List of dictionaries containing search results
+        """
+        print(f"🔍 Searching for: '{query_text}'")
+        
+        # Get query embedding
+        with torch.no_grad():
+            text_inputs = self.clip_processor(text=[query_text], padding=True, return_tensors="pt")
+            text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
+            
+            # Create a dummy image tensor to satisfy the model's requirements
+            dummy_image = torch.zeros(1, 3, 224, 224).to(self.device)
+            
+            outputs = self.main_model(**text_inputs, pixel_values=dummy_image)
+            query_embedding = outputs.text_embeds.cpu().numpy()
+        
+        # Calculate similarities
+        similarities = cosine_similarity(query_embedding, self.all_embeddings)[0]
+
+        # Get top-k results
+        top_indices = np.argsort(similarities)[::-1][:self.top_k * 2]  # Prendre plus de résultats
+        
+        results = []
+        for idx in top_indices:
+            if similarities[idx] > -0.5:
+                # Filter by category if specified
+                if filter_category and filter_category.lower() not in self.all_hierarchies[idx].lower():
+                    continue
+                    
+                results.append({
+                    'rank': len(results) + 1,
+                    'image_path': self.all_images[idx],
+                    'text': self.all_texts[idx],
+                    'color': self.all_colors[idx],
+                    'hierarchy': self.all_hierarchies[idx],
+                    'similarity': float(similarities[idx]),
+                    'index': int(idx),
+                    'url': self.all_urls[idx]
+                })
+                
+                if len(results) >= self.top_k:
+                    break
+        
+        print(f"✅ Found {len(results)} results")
+        return results
+    
+
+    def display_results(self, results: List[dict], query_info: str = ""):
+        """
+        Display search results with images and information
+        
+        Args:
+            results: List of search result dictionaries
+            query_info: Information about the query
+        """
+        if not results:
+            print("❌ No results found")
+            return
+        
+        print(f"\n🎯 Search Results for: {query_info}")
+        print("=" * 80)
+        
+        # Calculate grid layout
+        n_results = len(results)
+        cols = min(5, n_results)
+        rows = (n_results + cols - 1) // cols
+        
+        fig, axes = plt.subplots(rows, cols, figsize=(4*cols, 4*rows))
+        if rows == 1:
+            axes = axes.reshape(1, -1)
+        elif cols == 1:
+            axes = axes.reshape(-1, 1)
+        
+        for i, result in enumerate(results):
+            row = i // cols
+            col = i % cols
+            ax = axes[row, col]
+            
+            try:
+                # Load and display image
+                image = Image.open(result['image_path'])
+                ax.imshow(image)
+                ax.axis('off')
+                
+                # Add title with similarity score
+                title = f"#{result['rank']} (Similarity: {result['similarity']:.3f})\n{result['color']} {result['hierarchy']}"
+                ax.set_title(title, fontsize=10, wrap=True)
+                
+            except Exception as e:
+                ax.text(0.5, 0.5, f"Error loading image\n{result['image_path']}", 
+                       ha='center', va='center', transform=ax.transAxes)
+                ax.axis('off')
+        
+        # Hide empty subplots
+        for i in range(n_results, rows * cols):
+            row = i // cols
+            col = i % cols
+            axes[row, col].axis('off')
+        
+        plt.tight_layout()
+        plt.show()
+        
+        # Print detailed results
+        print("\n📋 Detailed Results:")
+        for result in results:
+            print(f"#{result['rank']:2d} | Similarity: {result['similarity']:.3f} | "
+                  f"Color: {result['color']:12s} | Category: {result['hierarchy']:15s} | "
+                  f"Text: {result['text'][:50]}...")
+            print(f"   🔗 URL: {result['url']}")
+            print()
+
+
+def main():
+    """Main function for command-line usage"""
+    parser = argparse.ArgumentParser(description="Fashion Search Engine with Category Emphasis")
+    parser.add_argument("--query", "-q", type=str, help="Search query")
+    parser.add_argument("--top-k", "-k", type=int, default=10, help="Number of results (default: 10)")
+    parser.add_argument("--fast", "-f", action="store_true", help="Fast mode (less items)")
+    parser.add_argument("--interactive", "-i", action="store_true", help="Interactive mode")
+
+    args = parser.parse_args()
+    
+    print("🎯 Fashion Search Engine with Category Emphasis")
+    
+    search_engine = FashionSearchEngine(
+        top_k=args.top_k, 
+    )
+    print("✅ Ready!")
+    
+    # Single query mode
+    if args.query:
+        print(f"🔍 Search: '{args.query}'...")
+        results = search_engine.search_by_text(args.query)
+        search_engine.display_results(results, args.query)
+
+    
+    # Interactive mode
+    print("Enter your query (e.g. 'red dress') or 'quit' to exit")
+    
+    while True:
+        try:
+            user_input = input("\n🔍 Query: ").strip()
+            if not user_input or user_input.lower() in ['quit', 'exit', 'q']:
+                print("👋 Goodbye!")
+                break
+            
+            if user_input.startswith('verify '):
+                if 'yellow accessories' in user_input:
+                    search_engine.display_yellow_accessories()
+                continue
+
+            print(f"🔍 Search: '{user_input}'...")
+            results = search_engine.search_by_text(user_input)
+            search_engine.display_results(results, user_input)
+            
+        except KeyboardInterrupt:
+            print("\n👋 Goodbye!")
+            break
+        except Exception as e:
+            print(f"❌ Error: {e}")
+
+if __name__ == "__main__":
+    main()