src/models\advanced_pretrained_classifier.py

"""
Advanced Pre-trained CNN Classifier for Architectural Style Classification
Uses multiple state-of-the-art architectures with ensemble methods.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from transformers import AutoImageProcessor, AutoModel
from typing import Dict, List, Tuple, Optional
import numpy as np


class AdvancedPretrainedClassifier(nn.Module):
    """
    Advanced pre-trained classifier using multiple architectures:
    - EfficientNetV2 (for general features)
    - ConvNeXt (for modern architectural features)
    - Swin Transformer (for hierarchical features)
    - Vision Transformer (for global attention)
    """
    
    def __init__(self, num_classes: int = 25, dropout_rate: float = 0.3):
        super().__init__()
        
        # Multiple pre-trained backbones
        self.efficientnet = timm.create_model(
            'tf_efficientnetv2_m', 
            pretrained=True, 
            num_classes=0,
            global_pool='avg'
        )
        
        self.convnext = timm.create_model(
            'convnext_base', 
            pretrained=True, 
            num_classes=0,
            global_pool='avg'
        )
        
        self.swin = timm.create_model(
            'swin_base_patch4_window7_224', 
            pretrained=True, 
            num_classes=0,
            global_pool='avg'
        )
        
        # Vision Transformer from HuggingFace
        self.vit_processor = AutoImageProcessor.from_pretrained('google/vit-base-patch16-224')
        self.vit = AutoModel.from_pretrained('google/vit-base-patch16-224')
        
        # Feature dimensions
        self.efficientnet_dim = self.efficientnet.num_features
        self.convnext_dim = self.convnext.num_features
        self.swin_dim = self.swin.num_features
        self.vit_dim = 768  # ViT base hidden size
        
        # Print feature dimensions for debugging
        print(f"Feature dimensions:")
        print(f"  EfficientNet: {self.efficientnet_dim}")
        print(f"  ConvNeXt: {self.convnext_dim}")
        print(f"  Swin: {self.swin_dim}")
        print(f"  ViT: {self.vit_dim}")
        
        # Feature fusion layers
        total_features = self.efficientnet_dim + self.convnext_dim + self.swin_dim + self.vit_dim
        
        self.feature_fusion = nn.Sequential(
            nn.Linear(total_features, 1024),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Dropout(dropout_rate)
        )
        
        # Multi-scale attention
        self.attention = MultiScaleAttention(
            efficientnet_dim=self.efficientnet_dim,
            convnext_dim=self.convnext_dim,
            swin_dim=self.swin_dim,
            vit_dim=self.vit_dim
        )
        
        # Final classifier with multiple heads
        self.classifier = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(256, num_classes)
        )
        
        # Auxiliary classifiers for each backbone
        self.aux_efficientnet = nn.Linear(self.efficientnet_dim, num_classes)
        self.aux_convnext = nn.Linear(self.convnext_dim, num_classes)
        self.aux_swin = nn.Linear(self.swin_dim, num_classes)
        self.aux_vit = nn.Linear(self.vit_dim, num_classes)
        
        # Temperature scaling for calibration
        self.temperature = nn.Parameter(torch.ones(1) * 1.5)
        
    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        # Extract features from each backbone
        efficientnet_features = self.efficientnet.forward_features(x)
        if isinstance(efficientnet_features, tuple):
            efficientnet_features = efficientnet_features[0]
        efficientnet_features = F.adaptive_avg_pool2d(efficientnet_features, 1).flatten(1)
        
        convnext_features = self.convnext.forward_features(x)
        if isinstance(convnext_features, tuple):
            convnext_features = convnext_features[0]
        convnext_features = F.adaptive_avg_pool2d(convnext_features, 1).flatten(1)
        
        swin_features = self.swin.forward_features(x)
        if isinstance(swin_features, tuple):
            swin_features = swin_features[0]
        swin_features = F.adaptive_avg_pool2d(swin_features, 1).flatten(1)
        
        # ViT features (need to process differently)
        vit_features = self._extract_vit_features(x)
        
        # Apply attention mechanism
        attended_features = self.attention(
            efficientnet_features, convnext_features, swin_features, vit_features
        )
        
        # Concatenate all features
        combined_features = torch.cat([
            efficientnet_features, convnext_features, swin_features, vit_features
        ], dim=1)
        
        # Feature fusion
        fused_features = self.feature_fusion(combined_features)
        
        # Main classifier
        main_logits = self.classifier(fused_features)
        
        # Auxiliary classifiers
        aux_efficientnet_logits = self.aux_efficientnet(efficientnet_features)
        aux_convnext_logits = self.aux_convnext(convnext_features)
        aux_swin_logits = self.aux_swin(swin_features)
        aux_vit_logits = self.aux_vit(vit_features)
        
        # Apply temperature scaling
        main_logits = main_logits / self.temperature
        
        return {
            'logits': main_logits,
            'aux_efficientnet': aux_efficientnet_logits,
            'aux_convnext': aux_convnext_logits,
            'aux_swin': aux_swin_logits,
            'aux_vit': aux_vit_logits,
            'features': fused_features,
            'attended_features': attended_features
        }
    
    def _extract_vit_features(self, x: torch.Tensor) -> torch.Tensor:
        """Extract features from Vision Transformer."""
        # Convert to PIL-like format for ViT
        # ViT expects normalized images in [0, 1] range
        x_normalized = x / 255.0
        
        # Use the CLS token output as features
        with torch.no_grad():
            outputs = self.vit(pixel_values=x_normalized)
            # Get the CLS token (first token)
            cls_output = outputs.last_hidden_state[:, 0, :]
        
        return cls_output


class MultiScaleAttention(nn.Module):
    """Multi-scale attention mechanism for feature fusion."""
    
    def __init__(self, efficientnet_dim: int, convnext_dim: int, swin_dim: int, vit_dim: int):
        super().__init__()
        
        # Project all features to a common dimension
        self.common_dim = 512
        
        # Projection layers to common dimension
        self.efficientnet_projection = nn.Linear(efficientnet_dim, self.common_dim)
        self.convnext_projection = nn.Linear(convnext_dim, self.common_dim)
        self.swin_projection = nn.Linear(swin_dim, self.common_dim)
        self.vit_projection = nn.Linear(vit_dim, self.common_dim)
        
        # Attention weights for each feature type
        self.efficientnet_attention = nn.Linear(self.common_dim, 1)
        self.convnext_attention = nn.Linear(self.common_dim, 1)
        self.swin_attention = nn.Linear(self.common_dim, 1)
        self.vit_attention = nn.Linear(self.common_dim, 1)
        
    def forward(self, efficientnet_features: torch.Tensor, convnext_features: torch.Tensor,
                swin_features: torch.Tensor, vit_features: torch.Tensor) -> torch.Tensor:
        
        # Project all features to common dimension
        efficientnet_proj = self.efficientnet_projection(efficientnet_features)
        convnext_proj = self.convnext_projection(convnext_features)
        swin_proj = self.swin_projection(swin_features)
        vit_proj = self.vit_projection(vit_features)
        
        # Calculate attention weights
        efficientnet_attn = torch.sigmoid(self.efficientnet_attention(efficientnet_proj))
        convnext_attn = torch.sigmoid(self.convnext_attention(convnext_proj))
        swin_attn = torch.sigmoid(self.swin_attention(swin_proj))
        vit_attn = torch.sigmoid(self.vit_attention(vit_proj))
        
        # Weighted features
        weighted_efficientnet = efficientnet_proj * efficientnet_attn
        weighted_convnext = convnext_proj * convnext_attn
        weighted_swin = swin_proj * swin_attn
        weighted_vit = vit_proj * vit_attn
        
        # Combine weighted features
        attended_features = (
            weighted_efficientnet + weighted_convnext + weighted_swin + weighted_vit
        ) / 4.0
        
        return attended_features


class AdvancedLossFunction(nn.Module):
    """Advanced loss function combining multiple loss types."""
    
    def __init__(self, num_classes: int = 25, alpha: float = 0.4, beta: float = 0.3, gamma: float = 0.3):
        super().__init__()
        self.alpha = alpha  # Main loss weight
        self.beta = beta    # Auxiliary loss weight
        self.gamma = gamma  # Focal loss weight
        
        # Loss functions
        self.cross_entropy = nn.CrossEntropyLoss(label_smoothing=0.1)
        self.focal_loss = FocalLoss(alpha=1.0, gamma=2.0)
        self.center_loss = CenterLoss(num_classes=num_classes, feat_dim=512)
        
    def forward(self, outputs: Dict[str, torch.Tensor], targets: torch.Tensor) -> Dict[str, torch.Tensor]:
        main_logits = outputs['logits']
        aux_logits = [
            outputs['aux_efficientnet'],
            outputs['aux_convnext'],
            outputs['aux_swin'],
            outputs['aux_vit']
        ]
        features = outputs['features']
        
        # Main classification loss
        main_loss = self.cross_entropy(main_logits, targets)
        
        # Auxiliary losses
        aux_losses = []
        for aux_logit in aux_logits:
            aux_loss = self.cross_entropy(aux_logit, targets)
            aux_losses.append(aux_loss)
        aux_loss = torch.mean(torch.stack(aux_losses))
        
        # Focal loss for hard examples
        focal_loss = self.focal_loss(main_logits, targets)
        
        # Center loss for feature learning
        center_loss = self.center_loss(features, targets)
        
        # Total loss
        total_loss = (
            self.alpha * main_loss +
            self.beta * aux_loss +
            self.gamma * focal_loss +
            0.1 * center_loss
        )
        
        return {
            'total_loss': total_loss,
            'main_loss': main_loss,
            'aux_loss': aux_loss,
            'focal_loss': focal_loss,
            'center_loss': center_loss
        }


class FocalLoss(nn.Module):
    """Focal Loss for handling class imbalance."""
    
    def __init__(self, alpha: float = 1.0, gamma: float = 2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        
    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()


class CenterLoss(nn.Module):
    """Center Loss for learning discriminative features."""
    
    def __init__(self, num_classes: int, feat_dim: int, device: str = 'cpu'):
        super().__init__()
        self.num_classes = num_classes
        self.feat_dim = feat_dim
        self.centers = nn.Parameter(torch.randn(num_classes, feat_dim))
        
    def forward(self, features: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        centers_batch = self.centers.index_select(0, targets)
        return F.mse_loss(features, centers_batch)


def create_advanced_classifier(num_classes: int = 25) -> AdvancedPretrainedClassifier:
    """Factory function to create the advanced classifier."""
    return AdvancedPretrainedClassifier(num_classes=num_classes)


def create_advanced_loss(num_classes: int = 25) -> AdvancedLossFunction:
    """Factory function to create the advanced loss function."""
    return AdvancedLossFunction(num_classes=num_classes)