new_model.py

import torch
import torch.nn as nn
from fightingcv_attention.attention.CBAM import CBAMBlock
from fightingcv_attention.attention.ECAAttention import ECAAttention
# from fightingcv_attention.attention.SEAttention import SEAttention
from fightingcv_attention.attention.BAM import BAMBlock
from fightingcv_attention.attention.ACmixAttention import ACmix
from torch.nn import init

# ================================
# Innovation Module 1: Global Context Embedding Module (FPA - Feature Pyramid Attention)
# ================================

class FPA(nn.Module):
    """
    Feature Pyramid Attention - Global Context Embedding Module
    Captures global context information through multi-scale feature pyramid and attention mechanism
    """
    def __init__(self, channels=2048):
        """
        Initialize FPA module
        Args:
            channels: Number of input feature channels, default 2048
        """
        super(FPA, self).__init__()
        channels_mid = int(channels/4)
        self.channels_cond = channels

        # Master branch - preserve original features
        self.conv_master = nn.Conv2d(self.channels_cond, channels, kernel_size=1, bias=False)
        self.bn_master = nn.BatchNorm2d(channels)

        # Global pooling branch - capture global information
        self.conv_gpb = nn.Conv2d(self.channels_cond, channels, kernel_size=1, bias=False)
        self.bn_gpb = nn.BatchNorm2d(channels)

        # Multi-scale feature extraction branch 1 (7x7 convolution)
        # C333 because of the shape of last feature maps is (16, 16).
        self.conv7x7_1 = nn.Conv2d(self.channels_cond, channels_mid, kernel_size=(7, 7), stride=2, padding=3, bias=False)
        self.bn1_1 = nn.BatchNorm2d(channels_mid)
        self.conv7x7_2 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(7, 7), stride=1, padding=3, bias=False)
        self.bn1_2 = nn.BatchNorm2d(channels_mid)

        # Multi-scale feature extraction branch 2 (5x5 convolution)
        self.conv5x5_1 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(5, 5), stride=2, padding=2, bias=False)
        self.bn2_1 = nn.BatchNorm2d(channels_mid)
        self.conv5x5_2 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(5, 5), stride=1, padding=2, bias=False)
        self.bn2_2 = nn.BatchNorm2d(channels_mid)

        # Multi-scale feature extraction branch 3 (3x3 convolution)
        self.conv3x3_1 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(3, 3), stride=2, padding=1, bias=False)
        self.bn3_1 = nn.BatchNorm2d(channels_mid)
        self.conv3x3_2 = nn.Conv2d(channels_mid, channels_mid, kernel_size=(3, 3), stride=1, padding=1, bias=False)
        self.bn3_2 = nn.BatchNorm2d(channels_mid)

        # Convolution Upsample - restore feature map size
        self.conv_upsample_3 = nn.ConvTranspose2d(channels_mid, channels_mid, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn_upsample_3 = nn.BatchNorm2d(channels_mid)
        self.conv_upsample_2 = nn.ConvTranspose2d(channels_mid, channels_mid, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn_upsample_2 = nn.BatchNorm2d(channels_mid)
        self.conv_upsample_1 = nn.ConvTranspose2d(channels_mid, channels, kernel_size=4, stride=2, padding=1, bias=False)
        self.bn_upsample_1 = nn.BatchNorm2d(channels)

        # Activation function and spatial attention
        self.relu = nn.ReLU(inplace=True)
        self.sp_attn = ECAAttention(kernel_size=5)

    def forward(self, x):
        """
        Forward propagation
        Args:
            x: Input feature maps [b, channels, h, w]
        Returns:
            out: Enhanced feature maps [b, channels, h, w]
        """
        # Master branch processing
        x_master = self.conv_master(x)
        x_master = self.bn_master(x_master)
        x_master = self.sp_attn(x_master)

        # Global pooling branch
        x_gpb = nn.AvgPool2d(x.shape[2:])(x).view(x.shape[0], self.channels_cond, 1, 1)
        x_gpb = self.conv_gpb(x_gpb)
        x_gpb = self.bn_gpb(x_gpb)

        # Branch 1 - Multi-scale feature extraction (7x7)
        x1_1 = self.conv7x7_1(x)
        x1_1 = self.bn1_1(x1_1)
        x1_1 = self.relu(x1_1)
        x1_2 = self.conv7x7_2(x1_1)
        x1_2 = self.bn1_2(x1_2)
        x1_2 = self.sp_attn(x1_2)

        # Branch 2 - Multi-scale feature extraction (5x5)
        x2_1 = self.conv5x5_1(x1_1)
        x2_1 = self.bn2_1(x2_1)
        x2_1 = self.relu(x2_1)
        x2_2 = self.conv5x5_2(x2_1)
        x2_2 = self.bn2_2(x2_2)
        x2_2 = self.sp_attn(x2_2)

        # Branch 3 - Multi-scale feature extraction (3x3)
        x3_1 = self.conv3x3_1(x2_1)
        x3_1 = self.bn3_1(x3_1)
        x3_1 = self.relu(x3_1)
        x3_2 = self.conv3x3_2(x3_1)
        x3_2 = self.bn3_2(x3_2)
        x3_2 = self.sp_attn(x3_2)

        # Merge branch 1 and 2 - bottom-up feature fusion
        x3_upsample = self.relu(self.bn_upsample_3(self.conv_upsample_3(x3_2)))
        x2_merge = self.relu(x2_2 + x3_upsample)
        x2_upsample = self.relu(self.bn_upsample_2(self.conv_upsample_2(x2_merge)))
        x1_merge = self.relu(x1_2 + x2_upsample)

        # Master branch fusion with multi-scale features
        x_master = x_master * self.relu(self.bn_upsample_1(self.conv_upsample_1(x1_merge)))

        # Final output: master branch + global branch
        out = self.relu(x_master + x_gpb)

        return out


# ================================
# Innovation Module 2: Channel AutoEncoder Module
# ================================

class Channel_AE(nn.Module):
    """
    Channel AutoEncoder Module
    Learns channel dependencies through encoder-decoder structure and fuses multi-layer features
    """
    def __init__(self, channel=512):
        """
        Initialize Channel AutoEncoder
        Args:
            channel: Number of input channels, default 512
        """
        super(Channel_AE, self).__init__()

        # Encoder part - progressively compress channel dimensions
        self.conv_1 = nn.Conv2d(channel, channel // 2, kernel_size=3, padding=1, bias=False)
        self.instance1 = nn.InstanceNorm2d(channel // 2)
        self.conv_2 = nn.Conv2d(channel // 2, channel // 4, kernel_size=3, padding=1, bias=False)
        self.instance2 = nn.InstanceNorm2d(channel // 4)
        self.conv_3 = nn.Conv2d(channel // 4, channel // 8, kernel_size=3, padding=1, bias=False)
        self.instance3 = nn.InstanceNorm2d(channel // 8)
        self.conv_4 = nn.Conv2d(channel // 8, channel // 16, kernel_size=3, padding=1, bias=False)
        self.instance4 = nn.InstanceNorm2d(channel // 16)
        self.conv_5 = nn.Conv2d(channel // 16, 1, kernel_size=3, padding=1, bias=False)

        # Decoder part - progressively restore channel dimensions
        self.de_conv_5 = nn.Conv2d(1, channel // 16, kernel_size=3, padding=1, bias=False)
        self.instance5 = nn.InstanceNorm2d(channel // 16)
        self.conv_6 = nn.Conv2d(channel // 16, channel // 8, kernel_size=3, padding=1, bias=False)
        self.instance6 = nn.InstanceNorm2d(channel // 8)
        self.conv_7 = nn.Conv2d(channel // 8, channel // 4, kernel_size=3, padding=1, bias=False)
        self.instance7 = nn.InstanceNorm2d(channel // 4)
        self.conv_8 = nn.Conv2d(channel // 4, channel // 2, kernel_size=3, padding=1, bias=False)
        self.instance8 = nn.InstanceNorm2d(channel // 2)
        self.conv_9 = nn.Conv2d(channel // 2, channel, kernel_size=3, padding=1, bias=False)
        self.instance9 = nn.InstanceNorm2d(channel)

        # Feature fusion layers
        self.de_conv_10 = nn.Conv2d(channel, channel//32, kernel_size=1, padding=0, bias=False)
        self.de_conv_0 = nn.Conv2d(1, channel//32, kernel_size=1, padding=0, bias=False)

        # Activation function
        self.relu = nn.ReLU(inplace=True)
        
        # CBAM attention mechanisms
        self.cbam1 = CBAMBlock(channel=channel//2, reduction=16, kernel_size=7)
        self.cbam2 = CBAMBlock(channel=channel//4, reduction=16, kernel_size=7)
        self.cbam3 = CBAMBlock(channel=channel//8, reduction=16, kernel_size=7)
        self.cbam4 = CBAMBlock(channel=channel//16, reduction=16, kernel_size=7)

    def init_weights(self):
        """Initialize weights for the module"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        """
        Forward propagation
        Args:
            x: Input feature maps [b, channel, h, w]
        Returns:
            out: Fused multi-layer features [b, channel//32 * 6, h, w]
        """
        # Encoder stage
        x = self.relu(self.instance1(self.conv_1(x)))
        x = self.cbam1(x)
        x = self.relu(self.instance2(self.conv_2(x)))
        x = self.cbam2(x)
        x = self.relu(self.instance3(self.conv_3(x)))
        x = self.cbam3(x)
        x = self.relu(self.instance4(self.conv_4(x)))
        x = self.cbam4(x)
        x = self.relu(self.conv_5(x))
        
        # Save encoder output
        x0 = x
        
        # Decoder stage
        x = self.relu(self.instance5(self.de_conv_5(x)))
        x = self.cbam4(x)
        x1 = x
        x = self.relu(self.instance6(self.conv_6(x)))
        x = self.cbam3(x)
        x2 = x
        x = self.relu(self.instance7(self.conv_7(x)))
        x = self.cbam2(x)
        x3 = x
        x = self.relu(self.instance8(self.conv_8(x)))
        x = self.cbam1(x)
        x4 = x
        x = self.relu(self.instance9(self.conv_9(x)))
        
        # Feature fusion
        x5 = self.de_conv_10(x)
        x0 = self.de_conv_0(x0)
        
        # Concatenate all layer features
        out = torch.cat([x5, x0, x1, x2, x3, x4], 1)
        
        return out


# ================================
# Multi-layer Channel AutoEncoder Application Module
# ================================

class DAM_layers(nn.Module):
    """
    Domain Adaptation Module Layers
    Apply Channel AutoEncoder to different layer feature maps
    """
    def __init__(self):
        super().__init__()
        self.dam_1 = Channel_AE(channel=256)   # Shallow layer features
        self.dam_2 = Channel_AE(channel=512)   # Middle layer features
        self.dam_3 = Channel_AE(channel=1024)  # Deep layer features

    def forward(self, x: list):
        """
        Forward propagation
        Args:
            x: List containing three different layer feature maps
        Returns:
            List of processed feature maps
        """
        x[0] = self.dam_1(x[0])
        x[1] = self.dam_2(x[1])
        x[2] = self.dam_3(x[2])
        return x


# ================================
# Additional Attention Modules
# ================================

class ChannelAttention(nn.Module):
    """Channel Attention Module"""
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.maxpool = nn.AdaptiveMaxPool2d(1)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.se = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(channel // reduction, channel, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result = self.maxpool(x)
        avg_result = self.avgpool(x)
        max_out = self.se(max_result)
        avg_out = self.se(avg_result)
        output = self.sigmoid(max_out + avg_out)
        return output


class SpatialAttention(nn.Module):
    """Spatial Attention Module"""
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result, _ = torch.max(x, dim=1, keepdim=True)
        avg_result = torch.mean(x, dim=1, keepdim=True)
        result = torch.cat([max_result, avg_result], 1)
        output = self.conv(result)
        output = self.sigmoid(output)
        return output


class CBAMBlock(nn.Module):
    """CBAM (Convolutional Block Attention Module)"""
    def __init__(self, channel=512, reduction=16, kernel_size=7):
        super().__init__()
        self.ca = ChannelAttention(channel=channel, reduction=reduction)
        self.sa = SpatialAttention(kernel_size=kernel_size)

    def init_weights(self):
        """Initialize weights for CBAM"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        b, c, _, _ = x.size()
        residual = x
        out = x * self.ca(x)
        out = out * self.sa(out)
        return out + residual

# ================================
# Usage Examples
# ================================

if __name__ == "__main__":
    # Test FPA module
    print("Testing Global Context Embedding Module (FPA):")
    fpa = FPA(channels=2048)
    x = torch.randn(2, 2048, 16, 16)
    out = fpa(x)
    print(f"Input shape: {x.shape}, Output shape: {out.shape}")
    
    # Test Channel_AE module
    print("\nTesting Channel AutoEncoder Module (Channel_AE):")
    channel_ae = Channel_AE(channel=512)
    x = torch.randn(2, 512, 32, 32)
    out = channel_ae(x)
    print(f"Input shape: {x.shape}, Output shape: {out.shape}")
    
    # Test DAM_layers module
    print("\nTesting Multi-layer Channel AutoEncoder Module (DAM_layers):")
    dam = DAM_layers()
    x_list = [
        torch.randn(2, 256, 64, 64),
        torch.randn(2, 512, 32, 32),
        torch.randn(2, 1024, 16, 16)
    ]
    out_list = dam(x_list)
    for i, (inp, outp) in enumerate(zip(x_list, out_list)):
        print(f"Layer {i+1} - Input shape: {inp.shape}, Output shape: {outp.shape}")