面向红外小目标检测的方向感知偏移增强与动态可变形卷积

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.ops import deform_conv2d


class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc1 = nn.Conv2d(in_planes, max(in_planes // 16, 1), kernel_size=1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(max(in_planes // 16, 1), in_planes, kernel_size=1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv1 = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)


class DeformableConv2d(nn.Module):
    # 带mask调制,即每个核的点都被0-1调制
    def __init__(self, in_dim, out_dim, kernel_size, stride=1, dilation=1, groups=1, bias=True, *,
                 offset_groups=1, with_mask=False
                 ):
        super().__init__()
        assert in_dim % groups == 0, "输入通道数必须能被groups整除"
        assert out_dim % groups == 0, "输出通道数必须能被groups整除"
        self.stride = stride
        self.padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
        self.dilation = dilation
        self.weight = nn.Parameter(torch.empty(out_dim, in_dim // groups, kernel_size, kernel_size))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_dim))
        else:
            self.bias = None
        nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='relu')
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_dim))
        else:
            self.bias = None

        self.with_mask = with_mask
        self.param_generator = nn.Conv2d(in_dim, (3 if with_mask else 2) * offset_groups * kernel_size ** 2, 3, 1, 1)
        nn.init.constant_(self.param_generator.weight, 0.0)  # 关键初始化
        nn.init.constant_(self.param_generator.bias, 0.0)  # 初始无偏移

    def forward(self, x):
        if self.with_mask:
            oh, ow, mask = self.param_generator(x).chunk(3, dim=1)
            offset = torch.cat([oh, ow], dim=1)
            mask = mask.sigmoid()
        else:
            offset = self.param_generator(x)
            mask = None

        x = deform_conv2d(
            x,
            offset=offset,
            weight=self.weight,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
            dilation=self.dilation,
            mask=mask,
        )
        return x


class MCCA(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16):
        super().__init__()
        self.reduction = reduction_ratio
        mid_channels = max(8, in_channels // reduction_ratio)  # 保证最小通道数

        # 水平分支（Width方向）
        self.width_pool = nn.AdaptiveAvgPool2d((1, None))  # [N,C,H,W] -> [N,C,1,W]
        self.width_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels, in_channels, 1),
            nn.Sigmoid()
        )

        # 垂直分支（Height方向）
        self.height_pool = nn.AdaptiveAvgPool2d((None, 1))  # [N,C,H,W] -> [N,C,H,1]
        self.height_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels, in_channels, 1),
            nn.Sigmoid()
        )

        # 全局分支（Global Average Pooling）
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.global_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels, in_channels, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        # 水平注意力
        w_pool = self.width_pool(x)  # [N,C,1,W]
        w_attn = self.width_conv(w_pool)  # [N,C,1,W]

        # 垂直注意力
        h_pool = self.height_pool(x)  # [N,C,H,1]
        h_attn = self.height_conv(h_pool)  # [N,C,H,1]

        # 全局注意力
        g_pool = self.global_pool(x)  # [N,C,1,1]
        g_attn = self.global_conv(g_pool)  # [N,C,1,1]

        # 合成三维注意力
        combined_attn = w_attn * h_attn * g_attn

        # 应用注意力权重
        return x * combined_attn


class MCDC(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super().__init__()
        self.kernel_size = kernel_size

        # Main convolution parameters
        self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels, kernel_size, kernel_size))

        # Offset and modulation generators
        self.offset_conv = nn.Conv2d(in_channels, 2 * kernel_size ** 2, kernel_size=3, padding=1)
        self.modulation_conv = nn.Conv2d(in_channels, kernel_size ** 2, kernel_size=3, padding=1)

        # MCCA module
        self.mcca = MCCA(in_channels)

        self.init_weights()

    def init_weights(self):
        nn.init.kaiming_uniform_(self.weight)
        nn.init.zeros_(self.offset_conv.weight)
        nn.init.zeros_(self.modulation_conv.weight)

    def forward(self, x):
        # Generate MCCA features
        x_attn = self.mcca(x)

        # Generate offsets and modulation
        offsets = self.offset_conv(x_attn)  # [N, 2*k*k, H, W]
        modulation = torch.sigmoid(self.modulation_conv(x_attn))  # [N, k*k, H, W]

        # Apply deformable convolution
        return deform_conv2d(
            input=x,
            offset=offsets,
            weight=self.weight,
            mask=modulation,
            padding=(self.kernel_size // 2, self.kernel_size // 2)
        )


class ChannelShuffle(nn.Module):
    def __init__(self, groups):
        super().__init__()
        self.groups = groups

    def forward(self, x):
        # 通道混洗实现
        batch_size, channels, height, width = x.size()
        channels_per_group = channels // self.groups
        x = x.view(batch_size, self.groups, channels_per_group, height, width)
        x = torch.transpose(x, 1, 2).contiguous()
        return x.view(batch_size, -1, height, width)


class LEFF(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        expanded_channels = in_channels * 4

        # 1x1卷积扩展通道
        self.conv1x1_expand = nn.Conv2d(in_channels, expanded_channels, kernel_size=1)

        # 深度卷积序列
        self.dw_conv3 = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=3,
                                  padding=1, groups=expanded_channels)
        self.dw_conv5 = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=5,
                                  padding=2, groups=expanded_channels)
        self.dw_conv3_2 = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=3,
                                    padding=1, groups=expanded_channels)

        # 通道混洗
        self.shuffle1 = ChannelShuffle(groups=expanded_channels // 2)
        self.shuffle2 = ChannelShuffle(groups=expanded_channels // 2)

        # 最后的1x1卷积恢复通道
        self.conv1x1_reduce = nn.Conv2d(expanded_channels, in_channels, kernel_size=1)

    def forward(self, x):
        # Step 1: 1x1卷积扩展通道
        x_1 = self.conv1x1_expand(x)  # [b, 4c, h, w]
        # Step 2: 3x3深度卷积
        x_2 = self.dw_conv3(x_1)  # [b, 4c, h, w]
        # Step 3: 通道混洗
        x2_s = self.shuffle1(x_2)  # [b, 4c, h, w]
        # Step 4: 5x5深度卷积
        x_3 = self.dw_conv5(x2_s)  # [b, 4c, h, w]
        # Step 5: 残差连接
        x_4 = x_3 + x_2  # [b, 4c, h, w]
        # Step 6: 通道混洗
        x4_s = self.shuffle2(x_4)  # [b, 4c, h, w]
        # Step 7: 3x5深度卷积
        x_5 = self.dw_conv3_2(x4_s)  # [b, 4c, h, w]
        # Step 8: 残差连接
        x_6 = x_5 + x_2  # [b, 4c, h, w]
        # Step 9: 恢复通道
        output = self.conv1x1_reduce(x_6)  # [b, c, h, w]

        return output


class DynamicConv(nn.Module):
    def __init__(self, in_ch, out_ch, dc_groups=3, offset_groups=3, kernel_list=[1, 3, 5], with_mask=True):
        super().__init__()
        self.num_convs = len(kernel_list)
        self.dc_groups = dc_groups
        self.offset_groups = offset_groups
        # 路由网络保持不变
        self.router = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_ch, self.num_convs, 1),
            nn.Softmax(dim=1)
        )
        self.with_mask = with_mask
        # 使用可变形卷积块替代普通卷积
        self.convs = nn.ModuleList([
            # nn.Conv2d(in_ch, out_ch, k, padding=k // 2)
            # MCDC(in_ch, out_ch, k)
            MCDC(in_ch, out_ch)
            # DeformableConv2d(in_dim=in_ch, out_dim=out_ch, kernel_size=k, groups=self.dc_groups,
            #                  offset_groups=self.offset_groups, with_mask=self.with_mask)
            for k in kernel_list
        ])

    def forward(self, x):
        B, _, H, W = x.shape
        weights = self.router(x).squeeze(-1).squeeze(-1)

        output = 0
        for i, conv in enumerate(self.convs):
            weight = weights[:, i].view(B, 1, 1, 1)  # 权重调整维度
            output += weight * conv(x)  # 加权求和

        return output


class Res_CBAM_block(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, dc_groups=4, offset_groups=4, kernel_list=[1, 3, 5],
                 with_mask=True):
        super(Res_CBAM_block, self).__init__()
        self.kernel_list = kernel_list
        self.with_mask = with_mask
        self.dc_groups = dc_groups
        self.offset_groups = offset_groups
        # self.conv1 =MCDC(in_channels, out_channels)
        # self.conv1 = DeformableConv2d(in_dim=in_channels, out_dim=out_channels, kernel_size=kernel_size, groups=self.dc_groups,
        #                      offset_groups=self.offset_groups, with_mask=self.with_mask)
        self.conv1 = DynamicConv(in_channels, out_channels, dc_groups=self.dc_groups, offset_groups=self.offset_groups,
                                 kernel_list=self.kernel_list, with_mask=with_mask)
        # self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        # self.conv2 = MCDC(out_channels, out_channels)
        # self.conv2 = DeformableConv2d(in_dim=out_channels, out_dim=out_channels, kernel_size=kernel_size,
        #                               groups=self.dc_groups,
        #                               offset_groups=self.offset_groups, with_mask=self.with_mask)
        self.conv2 = DynamicConv(out_channels, out_channels, dc_groups=self.dc_groups, offset_groups=self.offset_groups,
                                 kernel_list=self.kernel_list, with_mask=with_mask)
        # self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        if stride != 1 or out_channels != in_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
                nn.BatchNorm2d(out_channels))
        else:
            self.shortcut = None

        self.ca = ChannelAttention(out_channels)
        self.sa = SpatialAttention()

    def forward(self, x):
        residual = x
        if self.shortcut is not None:
            residual = self.shortcut(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.ca(out) * out
        out = self.sa(out) * out
        out += residual
        out = self.relu(out)
        return out


class Res_CBAM_block_org(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        if stride != 1 or out_channels != in_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
                nn.BatchNorm2d(out_channels))
        else:
            self.shortcut = None

        self.ca = ChannelAttention(out_channels)
        self.sa = SpatialAttention()

    def forward(self, x):
        residual = x
        if self.shortcut is not None:
            residual = self.shortcut(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.ca(out) * out
        out = self.sa(out) * out
        out += residual
        out = self.relu(out)
        return out


class UNet_dynamicDCNv2Block_MCCA_LEFF(nn.Module):
    def __init__(self, num_classes, input_channels, block, num_blocks, nb_filter, deep_supervision=True):
        super().__init__()
        self.relu = nn.ReLU(inplace=True)
        self.deep_supervision = deep_supervision
        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.down = nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=True)

        self.up_4 = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True)
        self.up_8 = nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True)
        self.up_16 = nn.Upsample(scale_factor=16, mode='bilinear', align_corners=True)

        self.conv0_0 = self._make_layer_org(Res_CBAM_block_org, input_channels, nb_filter[0])
        self.conv1_0 = self._make_layer(block, nb_filter[0], nb_filter[1], num_blocks[0])
        self.conv2_0 = self._make_layer(block, nb_filter[1], nb_filter[2], num_blocks[1])
        self.conv3_0 = self._make_layer(block, nb_filter[2], nb_filter[3], num_blocks[2])
        self.conv4_0 = self._make_layer(block, nb_filter[3], nb_filter[4], num_blocks[3])

        self.leff = LEFF(nb_filter[4])

        self.conv3_1 = self._make_layer(block, nb_filter[3] + nb_filter[4], nb_filter[3], num_blocks[2])
        self.conv2_2 = self._make_layer(block, nb_filter[2] + nb_filter[3], nb_filter[2], num_blocks[1])
        self.conv1_3 = self._make_layer(block, nb_filter[1] + nb_filter[2], nb_filter[1], num_blocks[0])
        self.conv0_4 = self._make_layer(block, nb_filter[0] + nb_filter[1], nb_filter[0])

        self.conv0_4_final = self._make_layer(block, nb_filter[0] * 5, nb_filter[0])

        self.conv0_4_1x1 = nn.Conv2d(nb_filter[4], nb_filter[0], kernel_size=1, stride=1)
        self.conv0_3_1x1 = nn.Conv2d(nb_filter[3], nb_filter[0], kernel_size=1, stride=1)
        self.conv0_2_1x1 = nn.Conv2d(nb_filter[2], nb_filter[0], kernel_size=1, stride=1)
        self.conv0_1_1x1 = nn.Conv2d(nb_filter[1], nb_filter[0], kernel_size=1, stride=1)
        if self.deep_supervision:
            self.final1 = nn.Conv2d(nb_filter[3], num_classes, kernel_size=1)
            self.final2 = nn.Conv2d(nb_filter[2], num_classes, kernel_size=1)
            self.final3 = nn.Conv2d(nb_filter[1], num_classes, kernel_size=1)
            self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
        else:
            self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)

    def _make_layer(self, block, input_channels, output_channels, num_blocks=1, dc_groups=4, offset_groups=4):
        layers = []
        layers.append(block(input_channels, output_channels, dc_groups=dc_groups, offset_groups=offset_groups))
        for i in range(num_blocks - 1):
            layers.append(block(output_channels, output_channels, dc_groups=dc_groups, offset_groups=offset_groups))
        return nn.Sequential(*layers)

    def _make_layer_org(self, block, input_channels, output_channels, num_blocks=1):
        layers = []
        layers.append(block(input_channels, output_channels))
        for i in range(num_blocks - 1):
            layers.append(block(output_channels, output_channels))
        return nn.Sequential(*layers)

    def forward(self, input):
        # 还差把input_x与最后的结果相加，然后输出结果
        x0_0 = self.conv0_0(input)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x2_0 = self.conv2_0(self.pool(x1_0))
        x3_0 = self.conv3_0(self.pool(x2_0))
        x4_0 = self.conv4_0(self.pool(x3_0))

        x4_0 = self.leff(x4_0)

        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))

        Final_x0_4 = self.conv0_4_final(
            torch.cat([self.up_16(self.conv0_4_1x1(x4_0)), self.up_8(self.conv0_3_1x1(x3_1)),
                       self.up_4(self.conv0_2_1x1(x2_2)), self.up(self.conv0_1_1x1(x1_3)), x0_4], 1))

        if self.deep_supervision:
            output1 = self.up_8(self.final1(x3_1))
            output2 = self.up_4(self.final2(x2_2))
            output3 = self.up(self.final3(x1_3))
            output4 = self.final4(Final_x0_4)
            return [output1, output2, output3, output4]
        else:
            output = self.final(Final_x0_4)
            return output


if __name__ == "__main__":
    inp = torch.ones((2, 1, 256, 256)).cuda()
    deep_supervision = "True"
    if deep_supervision == "True":
        model = UNet_dynamicDCNv2Block_MCCA_LEFF(1, 1, Res_CBAM_block, [2, 2, 2, 2], [16, 32, 64, 128, 256],
                                                 True).cuda()
        [output1, output2, output3, output4] = model(inp)
        print("output1", output1.shape)
        print("output2", output2.shape)
        print("output3", output3.shape)
        print("output4", output4.shape)
        '''
        output1 torch.Size([2, 1, 256, 256])
        output2 torch.Size([2, 1, 256, 256])
        output3 torch.Size([2, 1, 256, 256])
        output4 torch.Size([2, 1, 256, 256])
        '''
    else:
        model = UNet_dynamicDCNv2Block_MCCA_LEFF(1, 1, Res_CBAM_block, [2, 2, 2, 2], [16, 32, 64, 128, 256],
                                                 False).cuda()
        out = model(inp)
        print("out", out.shape)