工业缺陷检测实战——道路裂缝分割

AI街潜水的八角

于 2024-12-22 15:05:14 发布

阅读量877

点赞数 21

分类专栏：工业缺陷检测分割分类文章标签：机器学习深度学习人工智能

本文链接：https://blog.youkuaiyun.com/u013289254/article/details/144646528

版权

工业缺陷检测分割分类专栏收录该内容

9 篇文章

订阅专栏

第一步：准备数据

道路裂缝分割，总共118张

第二步：搭建模型

Attention Unet主要的中心思想就是提出来Attention gate模块，使用soft-attention替代hard-attention，将attention集成到Unet的跳跃连接和上采样模块中，实现空间上的注意力机制。通过attention机制来抑制图像中的无关信息，突出局部的重要特征

Attention Unet的模型结构和Unet十分相像，只是增加了Attention Gate模块来对skip connection和upsampling层做attention机制。

在Attention Gate模块中，g和xl分别为skip connection的输出和下一层的输出，如图

需要注意的是，在计算Wg和Wx后，对两者进行相加。但是，此时g的维度和xl的维度并不相等，则需要对g做下采样或对xl做上采样。（我倾向于对xl做上采样，因为在原本的Unet中，在Decoder就需要对下一层做上采样，所以，直接使用这个上采样结果可以减少网络计算）。

Wg和Wx经过相加，ReLU激活，1x1x1卷积，Sigmoid激活，生成一个权重信息，将这个权重与原始输入xl相乘，得到了对xl的attention激活。这就是Attenton Gate的思想。

Attenton Gate还有一个比较重要的特点是：这个权重可以经由网络学习！因为soft-attention是可微的，可以微分的attention就可以通过神经网络算出梯度并且前向传播和后向反馈来学习得到attention的权重。以此来学习更重要的特征。

第三步：代码

1）损失函数为：交叉熵损失函数+dice损失

2）网络代码：

class AttU_Net(nn.Module):
    def __init__(self,img_ch=3,output_ch=1):
        super(AttU_Net,self).__init__()
        
        self.Maxpool = nn.MaxPool2d(kernel_size=2,stride=2)

        self.Conv1 = conv_block(ch_in=img_ch,ch_out=64)
        self.Conv2 = conv_block(ch_in=64,ch_out=128)
        self.Conv3 = conv_block(ch_in=128,ch_out=256)
        self.Conv4 = conv_block(ch_in=256,ch_out=512)
        self.Conv5 = conv_block(ch_in=512,ch_out=1024)

        self.Up5 = up_conv(ch_in=1024,ch_out=512)
        self.Att5 = Attention_block(F_g=512,F_l=512,F_int=256)
        self.Up_conv5 = conv_block(ch_in=1024, ch_out=512)

        self.Up4 = up_conv(ch_in=512,ch_out=256)
        self.Att4 = Attention_block(F_g=256,F_l=256,F_int=128)
        self.Up_conv4 = conv_block(ch_in=512, ch_out=256)
        
        self.Up3 = up_conv(ch_in=256,ch_out=128)
        self.Att3 = Attention_block(F_g=128,F_l=128,F_int=64)
        self.Up_conv3 = conv_block(ch_in=256, ch_out=128)
        
        self.Up2 = up_conv(ch_in=128,ch_out=64)
        self.Att2 = Attention_block(F_g=64,F_l=64,F_int=32)
        self.Up_conv2 = conv_block(ch_in=128, ch_out=64)

        self.Conv_1x1 = nn.Conv2d(64,output_ch,kernel_size=1,stride=1,padding=0)


    def forward(self,x):
        # encoding path
        x1 = self.Conv1(x)

        x2 = self.Maxpool(x1)
        x2 = self.Conv2(x2)
        
        x3 = self.Maxpool(x2)
        x3 = self.Conv3(x3)

        x4 = self.Maxpool(x3)
        x4 = self.Conv4(x4)

        x5 = self.Maxpool(x4)
        x5 = self.Conv5(x5)

        # decoding + concat path
        d5 = self.Up5(x5)
        x4 = self.Att5(g=d5,x=x4)
        d5 = torch.cat((x4,d5),dim=1)        
        d5 = self.Up_conv5(d5)
        
        d4 = self.Up4(d5)
        x3 = self.Att4(g=d4,x=x3)
        d4 = torch.cat((x3,d4),dim=1)
        d4 = self.Up_conv4(d4)

        d3 = self.Up3(d4)
        x2 = self.Att3(g=d3,x=x2)
        d3 = torch.cat((x2,d3),dim=1)
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)
        x1 = self.Att2(g=d2,x=x1)
        d2 = torch.cat((x1,d2),dim=1)
        d2 = self.Up_conv2(d2)

        d1 = self.Conv_1x1(d2)

        return d1