extract_features.py

# *** USAGE ***
# > uv run python extract_features.py --class-name birthday_party

#!/usr/bin/env python3
"""
Image Feature Extraction Script

This script extracts high-quality image features using various state-of-the-art models:
- CLIP (OpenAI) - Original implementation
- DINOv2 (Meta) - Strong self-supervised features
- OpenCLIP (LAION) - Stronger CLIP variants
- ConvNeXt (Facebook) - Modern CNN features

Usage:
    python extract_features.py --model dinov2 --class-name cosmetic
    python extract_features.py --model openclip --class-name fashion
    python extract_features.py --model convnext --class-name electronics
    python extract_features.py --model clip --class-name cosmetic  # default
"""

import os
import torch
import json
import argparse
from pathlib import Path
from PIL import Image
import numpy as np
from transformers import CLIPProcessor, CLIPModel, AutoImageProcessor, AutoModel
from tqdm import tqdm
import warnings

# 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 ImageFeatureExtractor:
    """Extract features from images using various state-of-the-art models"""
    
    def __init__(self, model_type="clip", device=None):
        """
        Initialize the feature extractor
        
        Args:
            model_type (str): Type of model to use ('clip', 'dinov2', 'openclip', 'convnext')
            device (str): Device to run model on ('cuda', 'mps', 'cpu', or None for auto)
        """
        self.model_type = model_type
        self.device = self._get_device(device)
        
        # Model configurations
        self.model_configs = {
            "clip": {
                "model_name": "openai/clip-vit-base-patch32",
                "processor_class": CLIPProcessor,
                "model_class": CLIPModel,
                "feature_method": "get_image_features"
            },
            "openclip": {
                "model_name": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
                "processor_class": CLIPProcessor,
                "model_class": CLIPModel,
                "feature_method": "get_image_features"
            },
            "dinov2": {
                "model_name": "facebook/dinov2-base",
                "processor_class": AutoImageProcessor,
                "model_class": AutoModel,
                "feature_method": "pooler_output"
            },
            "convnext": {
                "model_name": "facebook/convnext-base-224-22k",
                "processor_class": AutoImageProcessor,
                "model_class": AutoModel,
                "feature_method": "pooler_output"
            }
        }
        
        if model_type not in self.model_configs:
            raise ValueError(f"Unsupported model type: {model_type}. Choose from: {list(self.model_configs.keys())}")
        
        config = self.model_configs[model_type]
        self.model_name = config["model_name"]
        
        print(f"🚀 Loading {model_type.upper()} model: {self.model_name}")
        print(f"🔧 Using device: {self.device}")
        
        # Load model and processor
        try:
            self.model = config["model_class"].from_pretrained(self.model_name, use_safetensors=True)
            self.processor = config["processor_class"].from_pretrained(self.model_name)
        except Exception as e:
            print(f"⚠️  Trying without safetensors...")
            self.model = config["model_class"].from_pretrained(self.model_name)
            self.processor = config["processor_class"].from_pretrained(self.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 _load_image(self, image_path):
        """Load and preprocess image"""
        try:
            image = Image.open(image_path)
            # Convert to RGB if needed (handles RGBA, grayscale, etc.)
            if image.mode != 'RGB':
                image = image.convert('RGB')
            return image
        except Exception as e:
            print(f"❌ Error loading {image_path}: {e}")
            return None
    
    def extract_single_feature(self, image_path):
        """Extract features from a single image"""
        image = self._load_image(image_path)
        if image is None:
            return None
        
        # Preprocess image
        inputs = self.processor(images=image, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Extract features based on model type
        with torch.no_grad():
            if self.model_type in ["clip", "openclip"]:
                image_features = self.model.get_image_features(**inputs)
            elif self.model_type in ["dinov2", "convnext"]:
                outputs = self.model(**inputs)
                if hasattr(outputs, 'pooler_output') and outputs.pooler_output is not None:
                    image_features = outputs.pooler_output
                elif hasattr(outputs, 'last_hidden_state'):
                    # Use global average pooling for models without pooler
                    image_features = outputs.last_hidden_state.mean(dim=1)
                else:
                    # Fallback to first token (CLS token)
                    image_features = outputs.last_hidden_state[:, 0]
            
            # Normalize features (important for similarity search)
            image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True)
        
        return image_features.cpu()  # Move back to CPU for saving
    
    def extract_batch_features(self, image_paths, batch_size=8):
        """Extract features from multiple images in batches"""
        features = []
        
        # Adjust batch size based on model type (larger models need smaller batches)
        if self.model_type in ["openclip", "dinov2"]:
            batch_size = max(1, batch_size // 2)
        
        for i in tqdm(range(0, len(image_paths), batch_size), desc="🔍 Extracting features"):
            batch_paths = image_paths[i:i + batch_size]
            batch_images = []
            valid_indices = []
            
            # Load batch of images
            for j, path in enumerate(batch_paths):
                image = self._load_image(path)
                if image is not None:
                    batch_images.append(image)
                    valid_indices.append(i + j)
                else:
                    features.append(None)  # Placeholder for failed image
            
            if not batch_images:
                continue
            
            try:
                # Process batch
                inputs = self.processor(images=batch_images, return_tensors="pt")
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                
                # Extract features based on model type
                with torch.no_grad():
                    if self.model_type in ["clip", "openclip"]:
                        batch_features = self.model.get_image_features(**inputs)
                    elif self.model_type in ["dinov2", "convnext"]:
                        outputs = self.model(**inputs)
                        if hasattr(outputs, 'pooler_output') and outputs.pooler_output is not None:
                            batch_features = outputs.pooler_output
                        elif hasattr(outputs, 'last_hidden_state'):
                            # Use global average pooling for models without pooler
                            batch_features = outputs.last_hidden_state.mean(dim=1)
                        else:
                            # Fallback to first token (CLS token)
                            batch_features = outputs.last_hidden_state[:, 0]
                    
                    # Normalize features
                    batch_features = batch_features / batch_features.norm(p=2, dim=-1, keepdim=True)
                
                # Add to results
                batch_features_cpu = batch_features.cpu()
                for j, feature in enumerate(batch_features_cpu):
                    features.append(feature.unsqueeze(0))  # Keep batch dimension
                    
            except Exception as e:
                print(f"❌ Error processing batch: {e}")
                # Add None for each image in failed batch
                for _ in batch_images:
                    features.append(None)
        
        return features
    
    def process_directory(self, images_dir="./images", embeddings_dir="./embeddings", 
                         class_name="default", batch_size=8):
        """
        Process all images in directory and save features
        
        Args:
            images_dir (str): Directory containing images
            embeddings_dir (str): Base directory to save features
            class_name (str): Class name for organizing embeddings (e.g., 'cosmetic', 'fashion')
            batch_size (int): Batch size for processing
        """
        images_path = Path(images_dir)
        
        # Create hierarchical path: embeddings_dir/class_name/model_type/
        embeddings_path = Path(embeddings_dir) / class_name / self.model_type
        
        # Create embeddings directory if it doesn't exist
        embeddings_path.mkdir(parents=True, exist_ok=True)
        
        # Find all image files
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
        image_files = [
            f for f in images_path.iterdir() 
            if f.is_file() and f.suffix.lower() in image_extensions
        ]
        
        if not image_files:
            print(f"❌ No images found in {images_dir}")
            return
        
        print(f"📁 Found {len(image_files)} images in {images_dir}")
        print(f"💾 Saving to: {embeddings_path}")
        
        # Extract features
        features = self.extract_batch_features(image_files, batch_size)
        
        # Save features and create mapping
        mapping = {}
        successful_extractions = 0
        
        print("💾 Saving features...")
        for image_file, feature in tqdm(zip(image_files, features), total=len(image_files), desc="Saving"):
            if feature is not None:
                # Create feature filename (same as image but with .pt extension)
                feature_filename = image_file.stem + ".pt"
                feature_path = embeddings_path / feature_filename
                
                # Save feature tensor
                torch.save(feature, feature_path)
                
                # Add to mapping
                mapping[str(image_file.name)] = feature_filename
                successful_extractions += 1
            else:
                print(f"⚠️  Skipped {image_file.name} due to processing error")
        
        # Save mapping file
        mapping_path = embeddings_path / "image_to_feature_mapping.json"
        with open(mapping_path, 'w') as f:
            json.dump(mapping, f, indent=2)
        
        # Save metadata
        metadata = {
            "class_name": class_name,
            "model_type": self.model_type,
            "model_name": self.model_name,
            "device_used": self.device,
            "total_images": len(image_files),
            "successful_extractions": successful_extractions,
            "failed_extractions": len(image_files) - successful_extractions,
            "feature_dimension": feature.shape[-1] if feature is not None else None,
            "batch_size": batch_size
        }
        
        metadata_path = embeddings_path / "extraction_metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)
        
        print(f"\n✅ Feature extraction complete!")
        print(f"📊 Successfully processed: {successful_extractions}/{len(image_files)} images")
        print(f"📁 Features saved to: {embeddings_path}")
        print(f"🗺️  Mapping saved to: {mapping_path}")
        print(f"📋 Metadata saved to: {metadata_path}")
        print(f"🏷️  Class: {class_name} | Model: {self.model_type}")
        
        if feature is not None:
            print(f"🔢 Feature dimension: {feature.shape[-1]}")


def main():
    """Main function to run feature extraction"""
    parser = argparse.ArgumentParser(description="Extract image features using various models")
    parser.add_argument("--model", type=str, default="clip", 
                        choices=["clip", "dinov2", "openclip", "convnext"],
                        help="Model type to use for feature extraction")
    parser.add_argument("--class-name", type=str, default="default",
                        help="Class name for organizing embeddings (e.g., 'cosmetic', 'fashion')")
    parser.add_argument("--batch-size", type=int, default=8,
                        help="Batch size for processing")
    parser.add_argument("--images-dir", type=str, default="./images",
                        help="Directory containing images")
    parser.add_argument("--embeddings-dir", type=str, default="./embeddings",
                        help="Base directory to save features")
    
    args = parser.parse_args()
    
    print(f"🎨 Image Feature Extraction with {args.model.upper()}")
    print(f"🏷️  Class: {args.class_name}")
    print("=" * 50)
    
    # Initialize extractor
    extractor = ImageFeatureExtractor(
        model_type=args.model,
        device=None  # Auto-detect best device
    )
    
    # Process images
    extractor.process_directory(
        images_dir=args.images_dir+"/"+args.class_name,
        embeddings_dir=args.embeddings_dir,
        class_name=args.class_name,
        batch_size=args.batch_size
    )


if __name__ == "__main__":
    main()