color_model.py

"""
ColorCLIP model for learning color-aligned embeddings.
This file contains the ColorCLIP model that learns to encode images and texts
in an embedding space specialized for color representation. It includes
a ResNet-based image encoder, a text encoder with custom tokenizer,
and contrastive loss functions for training.
"""

import config
import os
import json
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from PIL import Image
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd
from tqdm.auto import tqdm
from collections import defaultdict
from typing import Optional, List
import logging


# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# -------------------------------
# Dataset Classes
# -------------------------------
class ColorDataset(Dataset):
    """
    Dataset class for color embedding training.
    
    Handles loading images from local paths and tokenizing text descriptions
    for training the ColorCLIP model.
    """
    
    def __init__(self, dataframe, tokenizer, transform=None):
        """
        Initialize the color dataset.
        
        Args:
            dataframe: DataFrame with columns for image paths and text descriptions
            tokenizer: Tokenizer instance that converts text to list of integers (tokens)
            transform: Optional image transformations (default: standard ImageNet normalization)
        """
        self.df = dataframe.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.transform = transform or transforms.Compose([
            transforms.Resize((224,224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485,0.456,0.406],
                                 std=[0.229,0.224,0.225])
        ])
    
    def __len__(self):
        """Return the number of samples in the dataset."""
        return len(self.df)
    
    def __getitem__(self, idx):
        """
        Get a sample from the dataset.
        
        Args:
            idx: Index of the sample
            
        Returns:
            Tuple of (image_tensor, token_tensor)
        """
        row = self.df.iloc[idx]
        # Fix: Get the image path from the row, not the column name
        img_path = row[config.column_local_image_path]
        img = Image.open(img_path).convert("RGB")
        img = self.transform(img)
        tokens = torch.tensor(self.tokenizer(row[config.text_column]), dtype=torch.long)
        return img, tokens

# -------------------------------
# Tokenizer
# -------------------------------
class Tokenizer:
    """
    Tokenizer for extracting color-related keywords from text.
    
    This tokenizer filters text to keep only color-related words and basic
    descriptive words, then maps them to integer indices for embedding.
    """
    
    def __init__(self):
        """
        Initialize the tokenizer.
        
        Creates empty word-to-index and index-to-word mappings.
        Index 0 is reserved for padding/unknown tokens.
        """
        self.word2idx = defaultdict(lambda: 0)  # 0 = pad/unknown
        self.idx2word = {}
        self.counter = 1
    
    def preprocess_text(self, text):
        """
        Extract color-related keywords from text.
        
        Args:
            text: Input text string
            
        Returns:
            Preprocessed text containing only color and descriptive keywords
        """
        # Color-related keywords to keep
        color_keywords = ['red', 'blue', 'green', 'yellow', 'purple', 'pink', 'orange', 
                         'brown', 'black', 'white', 'gray', 'navy', 'beige', 'aqua', 'lime',
                         'violet', 'turquoise', 'teal', 'tan', 'snow', 'silver', 'plum',
                         'olive', 'fuchsia', 'gold', 'cream', 'ivory', 'maroon']
        
        # Keep only color-related words and basic descriptive words
        descriptive_words = ['shirt', 'dress', 'top', 'bottom', 'shoe', 'bag', 'hat', 'short', 'long', 'sleeve']
        
        words = text.lower().split()
        filtered_words = []
        for word in words:
            # Keep color words and some descriptive words
            if word in color_keywords or word in descriptive_words:
                filtered_words.append(word)
        
        return ' '.join(filtered_words) if filtered_words else text.lower()
    
    def fit(self, texts):
        """
        Build vocabulary from a list of texts.
        
        Args:
            texts: List of text strings to build vocabulary from
        """
        for text in texts:
            processed_text = self.preprocess_text(text)
            for word in processed_text.split():
                if word not in self.word2idx:
                    self.word2idx[word] = self.counter
                    self.idx2word[self.counter] = word
                    self.counter += 1
    
    def __call__(self, text):
        """
        Tokenize a text string into a list of integer indices.
        
        Args:
            text: Input text string
            
        Returns:
            List of integer token indices
        """
        processed_text = self.preprocess_text(text)
        return [self.word2idx[word] for word in processed_text.split()]
    
    def load_vocab(self, word2idx_dict):
        """
        Load vocabulary from a word-to-index dictionary.
        
        Args:
            word2idx_dict: Dictionary mapping words to indices
        """
        self.word2idx = defaultdict(lambda: 0, {k: int(v) for k, v in word2idx_dict.items()})
        self.idx2word = {int(v): k for k, v in word2idx_dict.items() if int(v) > 0}
        self.counter = max(self.word2idx.values(), default=0) + 1

# -------------------------------
# Model Components
# -------------------------------
class ImageEncoder(nn.Module):
    """
    Image encoder based on ResNet18 for extracting image embeddings.
    
    Uses a pretrained ResNet18 backbone and replaces the final layer
    to output embeddings of the specified dimension.
    """
    
    def __init__(self, embedding_dim=config.color_emb_dim):
        """
        Initialize the image encoder.
        
        Args:
            embedding_dim: Dimension of the output embedding (default: color_emb_dim)
        """
        super().__init__()
        self.backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.backbone.fc = nn.Sequential(
            nn.Dropout(0.1),  # Add regularization
            nn.Linear(self.backbone.fc.in_features, embedding_dim)
        )

    def forward(self, x):
        """
        Forward pass through the image encoder.
        
        Args:
            x: Image tensor [batch_size, channels, height, width]
            
        Returns:
            Normalized image embeddings [batch_size, embedding_dim]
        """
        x = self.backbone(x)
        return F.normalize(x, dim=-1)

class TextEncoder(nn.Module):
    """
    Text encoder for extracting text embeddings from token sequences.
    
    Uses an embedding layer followed by mean pooling (with optional length normalization)
    and a linear projection to the output embedding dimension.
    """
    
    def __init__(self, vocab_size, embedding_dim=config.color_emb_dim):
        """
        Initialize the text encoder.
        
        Args:
            vocab_size: Size of the vocabulary
            embedding_dim: Dimension of the output embedding (default: color_emb_dim)
        """
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, 32, padding_idx=0)  # Keep 32 dimensions
        self.dropout = nn.Dropout(0.1)  # Add regularization
        self.fc = nn.Linear(32, embedding_dim)

    def forward(self, x, lengths=None):
        """
        Forward pass through the text encoder.
        
        Args:
            x: Token tensor [batch_size, sequence_length]
            lengths: Optional sequence lengths tensor [batch_size] for proper mean pooling
            
        Returns:
            Normalized text embeddings [batch_size, embedding_dim]
        """
        emb = self.embedding(x)  # [B, T, 32]
        emb = self.dropout(emb)  # Apply dropout
        if lengths is not None:
            summed = emb.sum(dim=1)  # [B, 32]
            mean = summed / lengths.unsqueeze(1).clamp_min(1)
        else:
            mean = emb.mean(dim=1)
        return F.normalize(self.fc(mean), dim=-1)

class ColorCLIP(nn.Module):
    """
    Color CLIP model for learning color-aligned image-text embeddings.
    """
    def __init__(self, vocab_size, embedding_dim=config.color_emb_dim, tokenizer=None):
        """
        Initialize ColorCLIP model.
        
        Args:
            vocab_size: Size of the vocabulary for text encoding
            embedding_dim: Dimension of the embedding space (default: color_emb_dim)
            tokenizer: Optional Tokenizer instance (will create one if None)
        """
        super().__init__()
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.image_encoder = ImageEncoder(embedding_dim)
        self.text_encoder = TextEncoder(vocab_size, embedding_dim)
        self.tokenizer = tokenizer
    
    def forward(self, image, text, lengths=None):
        """
        Forward pass through the model.
        
        Args:
            image: Image tensor [B, C, H, W]
            text: Text token tensor [B, T]
            lengths: Optional sequence lengths tensor [B]
            
        Returns:
            Tuple of (image_embeddings, text_embeddings)
        """
        return self.image_encoder(image), self.text_encoder(text, lengths)

    def get_text_embeddings(self, texts: List[str]) -> torch.Tensor:
        """
        Get text embeddings for a list of text strings.
        
        Args:
            texts: List of text strings
            
        Returns:
            Text embeddings tensor [batch_size, embedding_dim]
        """
        if self.tokenizer is None:
            raise ValueError("Tokenizer must be set before calling get_text_embeddings")
            
        token_lists = [self.tokenizer(t) for t in texts]
        max_len = max((len(toks) for toks in token_lists), default=0)
        padded = [toks + [0] * (max_len - len(toks)) for toks in token_lists]
        input_ids = torch.tensor(padded, dtype=torch.long, device=next(self.parameters()).device)
        lengths = torch.tensor([len(toks) for toks in token_lists], dtype=torch.long, device=input_ids.device)
        with torch.no_grad():
            emb = self.text_encoder(input_ids, lengths)
        return emb
    
    @classmethod
    def from_pretrained(cls, model_path: str, vocab_path: Optional[str] = None, device: str = "cpu", repo_id: Optional[str] = None):
        """
        Load a pretrained ColorCLIP model from a file path or Hugging Face Hub.
        
        Args:
            model_path: Path to the model checkpoint (.pt file) or filename if using repo_id
            vocab_path: Optional path to tokenizer vocabulary JSON file or filename if using repo_id
            device: Device to load the model on (default: "cpu")
            repo_id: Optional Hugging Face repository ID (e.g., "username/model-name")
                     If provided, model_path and vocab_path should be filenames within the repo
            
        Returns:
            ColorCLIP model instance
            
        Example:
            # Load from local file
            model = ColorCLIP.from_pretrained("color_model.pt", "tokenizer_vocab.json")
            
            # Load from Hugging Face Hub
            from huggingface_hub import hf_hub_download
            model_file = hf_hub_download(repo_id="username/model-name", filename="color_model.pt")
            vocab_file = hf_hub_download(repo_id="username/model-name", filename="tokenizer_vocab.json")
            model = ColorCLIP.from_pretrained(model_file, vocab_file)
        """
        device_obj = torch.device(device)
        
        # Support loading from Hugging Face Hub if repo_id is provided
        if repo_id:
            try:
                from huggingface_hub import hf_hub_download
                model_path = hf_hub_download(repo_id=repo_id, filename=model_path)
                if vocab_path:
                    vocab_path = hf_hub_download(repo_id=repo_id, filename=vocab_path)
            except ImportError:
                raise ImportError("huggingface_hub is required to load models from Hugging Face. Install it with: pip install huggingface-hub")
        
        # Load model checkpoint
        checkpoint = torch.load(model_path, map_location=device_obj)
        
        # Extract vocab size and embedding dimension from checkpoint
        if isinstance(checkpoint, dict):
            # Try to get vocab_size from metadata first
            vocab_size = checkpoint.get('vocab_size', None)
            embedding_dim = checkpoint.get('embedding_dim', 16)
            
            # If not in metadata, try to infer from model state
            if vocab_size is None:
                state_dict = checkpoint.get('model_state_dict', checkpoint)
                if 'text_encoder.embedding.weight' in state_dict:
                    vocab_size = state_dict['text_encoder.embedding.weight'].shape[0]
                else:
                    raise ValueError("Could not determine vocab_size from checkpoint")
            
            # Load state dict
            state_dict = checkpoint.get('model_state_dict', checkpoint)
        else:
            raise ValueError("Checkpoint must be a dictionary")
        
        # Initialize model
        model = cls(vocab_size=vocab_size, embedding_dim=embedding_dim)
        model.load_state_dict(state_dict)
        model = model.to(device_obj)
        
        # Load tokenizer if vocab path is provided
        if vocab_path and os.path.exists(vocab_path):
            tokenizer = Tokenizer()
            with open(vocab_path, 'r') as f:
                vocab_dict = json.load(f)
            tokenizer.load_vocab(vocab_dict)
            model.tokenizer = tokenizer
        
        model.eval()
        return model
    
    def save_pretrained(self, save_directory: str, vocab_path: Optional[str] = None):
        """
        Save the model and optionally the tokenizer vocabulary.
        
        Args:
            save_directory: Directory to save the model
            vocab_path: Optional path to save tokenizer vocabulary
        """
        os.makedirs(save_directory, exist_ok=True)
        
        # Save model checkpoint
        model_path = os.path.join(save_directory, config.color_model_path)
        checkpoint = {
            'model_state_dict': self.state_dict(),
            'vocab_size': self.vocab_size,
            'embedding_dim': self.embedding_dim
        }
        torch.save(checkpoint, model_path)
        
        # Save tokenizer vocabulary if available
        if self.tokenizer is not None:
            vocab_dict = dict(self.tokenizer.word2idx)
            if vocab_path is None:
                vocab_path = os.path.join(save_directory, config.tokeniser_path)
            with open(vocab_path, 'w') as f:
                json.dump(vocab_dict, f)
        
        return model_path, vocab_path


# -------------------------------
# Loss Functions and Utilities
# -------------------------------
def clip_loss(image_emb, text_emb, temperature=0.07):
    """
    CLIP contrastive loss function.
    
    Args:
        image_emb: Image embeddings [batch_size, embedding_dim]
        text_emb: Text embeddings [batch_size, embedding_dim]
        temperature: Temperature scaling parameter
        
    Returns:
        Contrastive loss value
    """
    logits = image_emb @ text_emb.T / temperature
    labels = torch.arange(len(image_emb), device=image_emb.device)
    loss_i2t = F.cross_entropy(logits, labels)
    loss_t2i = F.cross_entropy(logits.T, labels)
    return (loss_i2t + loss_t2i) / 2

def collate_batch(batch):
    """
    Collate function for DataLoader that pads sequences and filters None values.
    
    Args:
        batch: List of (image, tokens) tuples or None
        
    Returns:
        Tuple of (images, padded_tokens, lengths) or None if batch is empty
    """
    batch = [b for b in batch if b is not None]
    if len(batch) == 0:
        return None
    imgs, tokens = zip(*batch)
    imgs = torch.stack(imgs, dim=0)
    lengths = torch.tensor([t.size(0) for t in tokens], dtype=torch.long)
    tokens_padded = nn.utils.rnn.pad_sequence(tokens, batch_first=True, padding_value=0)
    return imgs, tokens_padded, lengths



if __name__ == "__main__":
    """
    Training script for ColorCLIP model.
    This code only runs when the file is executed directly, not when imported.
    """
    # Configuration
    batch_size = 16
    lr = 1e-4
    epochs=50



    # Load dataset and split train/test 
    tokenizer = Tokenizer()
    df = pd.read_csv(config.local_dataset_path)
    
    # Data preparation: Reduce to main colors only (11 classes instead of 34)
    main_colors = ['beige', 'black', 'blue', 'brown', 'green', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
    df = df[df[config.color_column].isin(main_colors)].copy()
    print(f"📊 Filtered dataset: {len(df)} samples with {len(main_colors)} colors")
    print(f"🎨 Colors: {sorted(df[config.color_column].unique())}")
    
    tokenizer.fit(df[config.text_column].tolist())

    # Filter only rows with a valid local file
    df_local = df[df[config.column_local_image_path].astype(str).str.len() > 0]
    df_local = df_local[df_local[config.column_local_image_path].apply(lambda p: os.path.isfile(p))]
    df_local = df_local.reset_index(drop=True)


    # split 90/10
    df_local = df_local.sample(frac=1.0, random_state=42).reset_index(drop=True)
    split_idx = int(0.9 * len(df_local))
    df_train = df_local.iloc[:split_idx].reset_index(drop=True)
    df_test = df_local.iloc[split_idx:].reset_index(drop=True)


    train_dataset = ColorDataset(df_train, tokenizer)
    test_dataset = ColorDataset(df_test, tokenizer)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_batch, num_workers=0)
    test_loader  = DataLoader(test_dataset,  batch_size=batch_size, shuffle=False, collate_fn=collate_batch, num_workers=0)

    device = config.device
    print(f"Using device: {device}")

    model = ColorCLIP(vocab_size=tokenizer.counter, embedding_dim=config.color_emb_dim, tokenizer=tokenizer).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)  # Add weight decay

    # Save tokenizer vocab once (or update) so evaluation can reload the same mapping
    here = os.path.dirname(__file__)
    vocab_out = os.path.join(here, config.tokeniser_path)
    with open(vocab_out, "w") as f:
        json.dump(dict(tokenizer.word2idx), f)
    print(f"Tokenizer vocabulary saved to: {vocab_out}")


    for epoch in range(epochs):
        model.train()
        pbar = tqdm(total=len(train_loader), desc=f"Epoch {epoch+1}/{epochs} - train", leave=False)
        epoch_losses = []
        for batch in train_loader:
            if batch is None:
                pbar.update(1)
                continue
            imgs, texts, lengths = batch
            imgs = imgs.to(device)
            texts = texts.to(device)
            lengths = lengths.to(device)
            optimizer.zero_grad()
            img_emb, text_emb = model(imgs, texts, lengths)
            loss = clip_loss(img_emb, text_emb)
            loss.backward()
            optimizer.step()
            epoch_losses.append(loss.item())
            pbar.set_postfix({"loss": f"{loss.item():.4f}", "avg": f"{sum(epoch_losses)/len(epoch_losses):.4f}"})
            pbar.update(1)
        pbar.close()
        
        avg_train_loss = sum(epoch_losses) / len(epoch_losses) if epoch_losses else None
        if avg_train_loss is not None:
            print(f"[Train] Epoch {epoch+1}/{epochs} - avg loss: {avg_train_loss:.4f}")
        else:
            print(f"[Train] Epoch {epoch+1}/{epochs} - no valid batches")

        # Eval rapide sur test avec barre
        model.eval()
        test_losses = []
        with torch.no_grad():
            pbar_t = tqdm(total=len(test_loader), desc=f"Epoch {epoch+1}/{epochs} - test", leave=False)
            for batch in test_loader:
                if batch is None:
                    pbar_t.update(1)
                    continue
                imgs, texts, lengths = batch
                imgs = imgs.to(device)
                texts = texts.to(device)
                lengths = lengths.to(device)
                img_emb, text_emb = model(imgs, texts, lengths)
                test_losses.append(clip_loss(img_emb, text_emb).item())
                pbar_t.update(1)
            pbar_t.close()
        if len(test_losses) > 0:
            avg_test_loss = sum(test_losses) / len(test_losses)
            print(f"[Test ] Epoch {epoch+1}/{epochs} - avg loss: {avg_test_loss:.4f}")
        else:
            print(f"[Test ] Epoch {epoch+1}/{epochs} - no valid batches")

        # --- Save checkpoint at every epoch ---
        ckpt_dir = here
        latest_path = os.path.join(ckpt_dir, config.color_model_path)
        epoch_path = os.path.join(ckpt_dir, f"color_model_epoch_{epoch+1}.pt")
        checkpoint = {
            'model_state_dict': model.state_dict(),
            'vocab_size': model.vocab_size,
            'embedding_dim': model.embedding_dim
        }
        torch.save(checkpoint, latest_path)
        torch.save(checkpoint, epoch_path)
        print(f"[Save ] Saved checkpoints: {latest_path} and {epoch_path}")