CLANet中的CMIM模块代码复现

最新推荐文章于 2024-07-22 13:57:17 发布

yuride

最新推荐文章于 2024-07-22 13:57:17 发布

阅读量201

点赞数

文章标签：深度学习人工智能 python

本文链接：https://blog.csdn.net/weixin_43858206/article/details/128044222

版权

1、论文详情

论文名称：Collaborative Learning Attention Network Based
on RGB Image and Depth Image for Surface
Defect Inspection of No-Service Rail
文章链接：link

CMIM-n 模块细节图：
在这里插入图片描述

2、代码复现：

class CMIM(nn.Module):
    def __init__(self, in_channel):
        super(CMIM, self).__init__()
        self.GAP = nn.AdaptiveAvgPool1d(1)
        self.rgb_conv_1_1 = nn.Conv2d(in_channel, in_channel, 1)
        self.rgb_conv_1_2 = nn.Conv2d(in_channel, in_channel, 1)
        self.rgb_conv_2_1 = convsigmoid(in_channel, in_channel, 3)
        self.rgb_conv_2_2 = convsigmoid(in_channel, in_channel, 3)

        self.depth_conv_1_1 = nn.Conv2d(in_channel, in_channel, 1)
        self.depth_conv_1_2 = nn.Conv2d(in_channel, in_channel, 1)
        self.depth_conv_2_1 = convsigmoid(in_channel, in_channel, 3)
        self.depth_conv_2_2 = convsigmoid(in_channel, in_channel, 3)

        self.bn_rgb = nn.BatchNorm2d(in_channel)
        self.bn_depth = nn.BatchNorm2d(in_channel)
        self.relu = nn.ReLU(inplace=True)
        self.pool2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

    def forward(self, f_rgb, f_depth):
        b_1, c_1, h_1, w_1 = f_rgb.shape
        f_rgb_h = f_rgb.reshape(b_1 * c_1, h_1, w_1)
        f_rgb_w = f_rgb.reshape(b_1 * c_1, h_1, w_1).permute(0, 2, 1)
        x_gap_rgb = self.GAP(f_rgb_h).permute(0, 2, 1)
        y_gap_rgb = self.GAP(f_rgb_w).permute(0, 2, 1)
        z_hw_rgb = torch.cat((x_gap_rgb, y_gap_rgb), dim=2)
        z_hw_rgb = z_hw_rgb.reshape(b_1, c_1, 1, -1)
        z_hw_rgb = (self.bn_rgb(self.rgb_conv_1_1(z_hw_rgb))).mul(self.relu(self.bn_rgb(self.rgb_conv_1_2(z_hw_rgb))))
        z_hw_rgb = torch.split(z_hw_rgb, split_size_or_sections=h_1, dim=3)
        z_hw_h_rgb = z_hw_rgb[0]
        z_hw_w_rgb = z_hw_rgb[1]
        f_rgb_h_back = self.rgb_conv_2_1(z_hw_h_rgb)
        f_rgb_w_back = self.rgb_conv_2_2(z_hw_w_rgb)
        f_rgb_multip = f_rgb_h_back.mul(f_rgb_w_back)
        stage_1 = f_rgb_multip.mul(f_depth)

        b_2, c_2, h_2, w_2 = stage_1.shape
        f_depth_h = stage_1.reshape(b_2 * c_2, h_2, w_2)
        f_depth_w = stage_1.reshape(b_2 * c_2, h_2, w_2).permute(0, 2, 1)
        x_gap_depth = self.GAP(f_depth_h).permute(0, 2, 1)
        y_gap_depth = self.GAP(f_depth_w).permute(0, 2, 1)
        z_hw_depth = torch.cat((x_gap_depth, y_gap_depth), dim=2)
        z_hw_depth = z_hw_depth.reshape(b_2, c_2, 1, -1)
        z_hw_depth = (self.bn_depth(self.depth_conv_1_1(z_hw_depth))).mul(self.relu(self.bn_depth(self.depth_conv_1_2(z_hw_depth))))
        z_hw_depth = torch.split(z_hw_depth, split_size_or_sections=h_2, dim=3)
        z_hw_h_depth = z_hw_depth[0]
        z_hw_w_depth = z_hw_depth[1]
        f_depth_h_back = self.depth_conv_2_1(z_hw_h_depth)
        f_depth_w_back = self.depth_conv_2_2(z_hw_w_depth)
        f_depth_multip = f_depth_h_back.mul(f_depth_w_back)
        out = f_depth_multip.mul(f_rgb)

        return out