model_architecture.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
import numpy as np

dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier",
              "Appenzeller", "Australian_Terrier", "Bedlington_Terrier", "Bernese_Mountain_Dog", "Bichon_Frise",
              "Blenheim_Spaniel", "Border_Collie", "Border_Terrier", "Boston_Bull", "Bouvier_Des_Flandres",
              "Brabancon_Griffon", "Brittany_Spaniel", "Cardigan", "Chesapeake_Bay_Retriever",
              "Chihuahua", "Dachshund", "Dandie_Dinmont", "Doberman", "English_Foxhound", "English_Setter",
              "English_Springer", "EntleBucher", "Eskimo_Dog", "French_Bulldog", "German_Shepherd",
              "German_Short-Haired_Pointer", "Gordon_Setter", "Great_Dane", "Great_Pyrenees",
              "Greater_Swiss_Mountain_Dog","Havanese", "Ibizan_Hound", "Irish_Setter", "Irish_Terrier",
              "Irish_Water_Spaniel", "Irish_Wolfhound", "Italian_Greyhound", "Japanese_Spaniel",
              "Kerry_Blue_Terrier", "Labrador_Retriever", "Lakeland_Terrier", "Leonberg", "Lhasa",
              "Maltese_Dog", "Mexican_Hairless", "Newfoundland", "Norfolk_Terrier", "Norwegian_Elkhound",
              "Norwich_Terrier", "Old_English_Sheepdog", "Pekinese", "Pembroke", "Pomeranian",
              "Rhodesian_Ridgeback", "Rottweiler", "Saint_Bernard", "Saluki", "Samoyed",
              "Scotch_Terrier", "Scottish_Deerhound", "Sealyham_Terrier", "Shetland_Sheepdog", "Shiba_Inu",
              "Shih-Tzu", "Siberian_Husky", "Staffordshire_Bullterrier", "Sussex_Spaniel",
              "Tibetan_Mastiff", "Tibetan_Terrier", "Walker_Hound", "Weimaraner",
              "Welsh_Springer_Spaniel", "West_Highland_White_Terrier", "Yorkshire_Terrier",
              "Affenpinscher", "Basenji", "Basset", "Beagle", "Black-and-Tan_Coonhound", "Bloodhound",
              "Bluetick", "Borzoi", "Boxer", "Briard", "Bull_Mastiff", "Cairn", "Chow", "Clumber",
              "Cocker_Spaniel", "Collie", "Curly-Coated_Retriever", "Dhole", "Dingo",
              "Flat-Coated_Retriever", "Giant_Schnauzer", "Golden_Retriever", "Groenendael", "Keeshond",
              "Kelpie", "Komondor", "Kuvasz", "Malamute", "Malinois", "Miniature_Pinscher",
              "Miniature_Poodle", "Miniature_Schnauzer", "Otterhound", "Papillon", "Pug", "Redbone",
              "Schipperke", "Silky_Terrier", "Soft-Coated_Wheaten_Terrier", "Standard_Poodle",
              "Standard_Schnauzer", "Toy_Poodle", "Toy_Terrier", "Vizsla", "Whippet",
              "Wire-Haired_Fox_Terrier"]


class MorphologicalFeatureExtractor(nn.Module):

    def __init__(self, in_features):
        super().__init__()

        # 基礎特徵維度設置
        self.reduced_dim = in_features // 4
        self.spatial_size = max(7, int(np.sqrt(self.reduced_dim // 64)))

        # 1. 特徵空間轉換器：將一維特徵轉換為二維空間表示
        self.dimension_transformer = nn.Sequential(
            nn.Linear(in_features, self.spatial_size * self.spatial_size * 64),
            nn.LayerNorm(self.spatial_size * self.spatial_size * 64),
            nn.ReLU()
        )

        # 2. 形態特徵分析器：分析具體的形態特徵
        self.morphological_analyzers = nn.ModuleDict({
            # 體型分析器：分析整體比例和大小
            'body_proportion': nn.Sequential(
                # 使用大卷積核捕捉整體體型特徵
                nn.Conv2d(64, 128, kernel_size=7, padding=3),
                nn.BatchNorm2d(128),
                nn.ReLU(),
                # 使用較小的卷積核精煉特徵
                nn.Conv2d(128, 128, kernel_size=3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU()
            ),

            # 頭部特徵分析器：關注耳朵、臉部等
            'head_features': nn.Sequential(
                # 中等大小的卷積核，適合分析頭部結構
                nn.Conv2d(64, 128, kernel_size=5, padding=2),
                nn.BatchNorm2d(128),
                nn.ReLU(),
                # 小卷積核捕捉細節
                nn.Conv2d(128, 128, kernel_size=3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU()
            ),

            # 尾部特徵分析器
            'tail_features': nn.Sequential(
                nn.Conv2d(64, 128, kernel_size=5, padding=2),
                nn.BatchNorm2d(128),
                nn.ReLU(),
                nn.Conv2d(128, 128, kernel_size=3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU()
            ),

            # 毛髮特徵分析器：分析毛髮長度、質地等
            'fur_features': nn.Sequential(
                # 使用多個小卷積核捕捉毛髮紋理
                nn.Conv2d(64, 128, kernel_size=3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU(),
                nn.Conv2d(128, 128, kernel_size=3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU()
            ),

            # 顏色特徵分析器：分析顏色分佈
            'color_pattern': nn.Sequential(
                # 第一層：捕捉基本顏色分布
                nn.Conv2d(64, 128, kernel_size=3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU(),

                # 第二層：分析顏色模式和花紋
                nn.Conv2d(128, 128, kernel_size=3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU(),

                # 第三層：整合顏色信息
                nn.Conv2d(128, 128, kernel_size=1),
                nn.BatchNorm2d(128),
                nn.ReLU()
            )
        })

        # 3. 特徵注意力機制：動態關注不同特徵
        self.feature_attention = nn.MultiheadAttention(
            embed_dim=128,
            num_heads=8,
            dropout=0.1,
            batch_first=True
        )

        # 4. 特徵關係分析器：分析不同特徵之間的關係
        self.relation_analyzer = nn.Sequential(
            nn.Linear(128 * 5, 256),  # 4個特徵分析器的輸出
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.LayerNorm(128),
            nn.ReLU()
        )

        # 5. 特徵整合器：將所有特徵智能地組合在一起
        self.feature_integrator = nn.Sequential(
            nn.Linear(128 * 6, in_features),  # 5個原始特徵 + 1個關係特徵
            nn.LayerNorm(in_features),
            nn.ReLU()
        )

    def forward(self, x):
        batch_size = x.size(0)

        # 1. 將特徵轉換為空間形式
        spatial_features = self.dimension_transformer(x).view(
            batch_size, 64, self.spatial_size, self.spatial_size
        )

        # 2. 分析各種形態特徵
        morphological_features = {}
        for name, analyzer in self.morphological_analyzers.items():
            # 提取特定形態特徵
            features = analyzer(spatial_features)
            # 使用自適應池化統一特徵大小
            pooled_features = F.adaptive_avg_pool2d(features, (1, 1))
            # 重塑特徵為向量形式
            morphological_features[name] = pooled_features.view(batch_size, -1)

        # 3. 特徵注意力處理
        # 將所有特徵堆疊成序列
        stacked_features = torch.stack(list(morphological_features.values()), dim=1)
        # 應用注意力機制
        attended_features, _ = self.feature_attention(
            stacked_features, stacked_features, stacked_features
        )

        # 4. 分析特徵之間的關係
        # 將所有特徵連接起來
        combined_features = torch.cat(list(morphological_features.values()), dim=1)
        # 提取特徵間的關係
        relation_features = self.relation_analyzer(combined_features)

        # 5. 特徵整合
        # 將原始特徵和關係特徵結合
        final_features = torch.cat([
            *morphological_features.values(),
            relation_features
        ], dim=1)

        # 6. 最終整合
        integrated_features = self.feature_integrator(final_features)

        # 添加殘差連接
        return integrated_features + x


class MultiHeadAttention(nn.Module):

    def __init__(self, in_dim, num_heads=8):
        """
        Initializes the MultiHeadAttention module.
        Args:
            in_dim (int): Dimension of the input features.
            num_heads (int): Number of attention heads. Defaults to 8.
        """
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = max(1, in_dim // num_heads)  
        self.scaled_dim = self.head_dim * num_heads  
        self.fc_in = nn.Linear(in_dim, self.scaled_dim)  
        self.query = nn.Linear(self.scaled_dim, self.scaled_dim)  # Query projection
        self.key = nn.Linear(self.scaled_dim, self.scaled_dim)  # Key projection
        self.value = nn.Linear(self.scaled_dim, self.scaled_dim)  # Value projection
        self.fc_out = nn.Linear(self.scaled_dim, in_dim)  # Linear layer to project output back to in_dim

    def forward(self, x):
        """
        Forward pass for multi-head attention mechanism.
        Args:
            x (Tensor): Input tensor of shape (batch_size, input_dim).
            x 是 (N,D), N:批次大小, D:輸入特徵維度
        Returns:
            Tensor: Output tensor after applying attention mechanism.
        """
        N = x.shape[0]  # Batch size
        x = self.fc_in(x)  # Project input to scaled_dim
        q = self.query(x).view(N, self.num_heads, self.head_dim)  # Compute queries
        k = self.key(x).view(N, self.num_heads, self.head_dim)  # Compute keys
        v = self.value(x).view(N, self.num_heads, self.head_dim)  # Compute values

        # Calculate attention scores
        energy = torch.einsum("nqd,nkd->nqk", [q, k])  
        attention = F.softmax(energy / (self.head_dim ** 0.5), dim=2)  # Apply softmax with scaling

        # Compute weighted sum of values based on attention scores
        out = torch.einsum("nqk,nvd->nqd", [attention, v]) 
        out = out.reshape(N, self.scaled_dim)  # Concatenate all heads
        out = self.fc_out(out)  # Project back to original input dimension
        return out


class BaseModel(nn.Module):

    def __init__(self, num_classes, device='cuda' if torch.cuda.is_available() else 'cpu'):
        super().__init__()
        self.device = device

        # 1. Initialize backbone
        self.backbone = timm.create_model(
                'convnextv2_base',
                pretrained=True,
                num_classes=0
        )

        # 2. 使用測試數據來確定實際的特徵維度
        with torch.no_grad():  
            dummy_input = torch.randn(1, 3, 224, 224)
            features = self.backbone(dummy_input)

            if len(features.shape) > 2:
                features = features.mean([-2, -1])

            self.feature_dim = features.shape[1]

        print(f"Feature Dimension from V2 backbone: {self.feature_dim}")

        # 3. Setup multi-head attention layer
        self.num_heads = max(1, min(8, self.feature_dim // 64))
        self.attention = MultiHeadAttention(self.feature_dim, num_heads=self.num_heads)

        # 4. Setup classifier
        self.classifier = nn.Sequential(
            nn.LayerNorm(self.feature_dim),
            nn.Dropout(0.3),
            nn.Linear(self.feature_dim, num_classes)
        )

        self.morphological_extractor = MorphologicalFeatureExtractor(
            in_features=self.feature_dim
        )

        self.feature_fusion = nn.Sequential(
            nn.Linear(self.feature_dim * 3, self.feature_dim),  
            nn.LayerNorm(self.feature_dim),
            nn.ReLU(),
            nn.Linear(self.feature_dim, self.feature_dim),
            nn.LayerNorm(self.feature_dim),
            nn.ReLU()
        )

    def forward(self, x):
        """
        Forward propagation process, combining V2's FCCA and multi-head attention mechanism
        Args:
            x (Tensor): Input image tensor of shape [batch_size, channels, height, width]
        Returns:
            Tuple[Tensor, Tensor]: Classification logits and attention features
        """
        x = x.to(self.device)

        # 1. Extract base features
        features = self.backbone(x)
        if len(features.shape) > 2:
            features = features.mean([-2, -1])

        # 2. Extract morphological features (including all detail features)
        morphological_features = self.morphological_extractor(features)

        # 3. Feature fusion (note dimension alignment with new fusion layer)
        combined_features = torch.cat([
            features,  # Original features
            morphological_features,  # Morphological features
            features * morphological_features  # Feature interaction information
        ], dim=1)
        fused_features = self.feature_fusion(combined_features)

        # 4. Apply attention mechanism
        attended_features = self.attention(fused_features)

        # 5. Final classifier
        logits = self.classifier(attended_features)

        return logits, attended_features