| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189 |
- # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- import paddle
- import paddle.nn as nn
- from paddlex.paddleseg.cvlibs import manager
- from paddlex.paddleseg.models import layers
- from paddlex.paddleseg import utils
- import numpy as np
- @manager.MODELS.add_component
- class AttentionUNet(nn.Layer):
- """
- The Attention-UNet implementation based on PaddlePaddle.
- As mentioned in the original paper, author proposes a novel attention gate (AG)
- that automatically learns to focus on target structures of varying shapes and sizes.
- Models trained with AGs implicitly learn to suppress irrelevant regions in an input image while
- highlighting salient features useful for a specific task.
- The original article refers to
- Oktay, O, et, al. "Attention u-net: Learning where to look for the pancreas."
- (https://arxiv.org/pdf/1804.03999.pdf).
- Args:
- num_classes (int): The unique number of target classes.
- pretrained (str, optional): The path or url of pretrained model. Default: None.
- """
- def __init__(self, num_classes, pretrained=None):
- super().__init__()
- n_channels = 3
- self.encoder = Encoder(n_channels, [64, 128, 256, 512])
- filters = np.array([64, 128, 256, 512, 1024])
- self.up5 = UpConv(ch_in=filters[4], ch_out=filters[3])
- self.att5 = AttentionBlock(
- F_g=filters[3], F_l=filters[3], F_out=filters[2])
- self.up_conv5 = ConvBlock(ch_in=filters[4], ch_out=filters[3])
- self.up4 = UpConv(ch_in=filters[3], ch_out=filters[2])
- self.att4 = AttentionBlock(
- F_g=filters[2], F_l=filters[2], F_out=filters[1])
- self.up_conv4 = ConvBlock(ch_in=filters[3], ch_out=filters[2])
- self.up3 = UpConv(ch_in=filters[2], ch_out=filters[1])
- self.att3 = AttentionBlock(
- F_g=filters[1], F_l=filters[1], F_out=filters[0])
- self.up_conv3 = ConvBlock(ch_in=filters[2], ch_out=filters[1])
- self.up2 = UpConv(ch_in=filters[1], ch_out=filters[0])
- self.att2 = AttentionBlock(
- F_g=filters[0], F_l=filters[0], F_out=filters[0] // 2)
- self.up_conv2 = ConvBlock(ch_in=filters[1], ch_out=filters[0])
- self.conv_1x1 = nn.Conv2D(
- filters[0], num_classes, kernel_size=1, stride=1, padding=0)
- self.pretrained = pretrained
- self.init_weight()
- def forward(self, x):
- x5, (x1, x2, x3, x4) = self.encoder(x)
- d5 = self.up5(x5)
- x4 = self.att5(g=d5, x=x4)
- d5 = paddle.concat([x4, d5], axis=1)
- d5 = self.up_conv5(d5)
- d4 = self.up4(d5)
- x3 = self.att4(g=d4, x=x3)
- d4 = paddle.concat((x3, d4), axis=1)
- d4 = self.up_conv4(d4)
- d3 = self.up3(d4)
- x2 = self.att3(g=d3, x=x2)
- d3 = paddle.concat((x2, d3), axis=1)
- d3 = self.up_conv3(d3)
- d2 = self.up2(d3)
- x1 = self.att2(g=d2, x=x1)
- d2 = paddle.concat((x1, d2), axis=1)
- d2 = self.up_conv2(d2)
- logit = self.conv_1x1(d2)
- logit_list = [logit]
- return logit_list
- def init_weight(self):
- if self.pretrained is not None:
- utils.load_entire_model(self, self.pretrained)
- class AttentionBlock(nn.Layer):
- def __init__(self, F_g, F_l, F_out):
- super().__init__()
- self.W_g = nn.Sequential(
- nn.Conv2D(
- F_g, F_out, kernel_size=1, stride=1, padding=0),
- nn.BatchNorm2D(F_out))
- self.W_x = nn.Sequential(
- nn.Conv2D(
- F_l, F_out, kernel_size=1, stride=1, padding=0),
- nn.BatchNorm2D(F_out))
- self.psi = nn.Sequential(
- nn.Conv2D(
- F_out, 1, kernel_size=1, stride=1, padding=0),
- nn.BatchNorm2D(1),
- nn.Sigmoid())
- self.relu = nn.ReLU()
- def forward(self, g, x):
- g1 = self.W_g(g)
- x1 = self.W_x(x)
- psi = self.relu(g1 + x1)
- psi = self.psi(psi)
- res = x * psi
- return res
- class UpConv(nn.Layer):
- def __init__(self, ch_in, ch_out):
- super().__init__()
- self.up = nn.Sequential(
- nn.Upsample(
- scale_factor=2, mode="bilinear"),
- nn.Conv2D(
- ch_in, ch_out, kernel_size=3, stride=1, padding=1),
- nn.BatchNorm2D(ch_out),
- nn.ReLU())
- def forward(self, x):
- return self.up(x)
- class Encoder(nn.Layer):
- def __init__(self, input_channels, filters):
- super().__init__()
- self.double_conv = nn.Sequential(
- layers.ConvBNReLU(input_channels, 64, 3),
- layers.ConvBNReLU(64, 64, 3))
- down_channels = filters
- self.down_sample_list = nn.LayerList([
- self.down_sampling(channel, channel * 2)
- for channel in down_channels
- ])
- def down_sampling(self, in_channels, out_channels):
- modules = []
- modules.append(nn.MaxPool2D(kernel_size=2, stride=2))
- modules.append(layers.ConvBNReLU(in_channels, out_channels, 3))
- modules.append(layers.ConvBNReLU(out_channels, out_channels, 3))
- return nn.Sequential(*modules)
- def forward(self, x):
- short_cuts = []
- x = self.double_conv(x)
- for down_sample in self.down_sample_list:
- short_cuts.append(x)
- x = down_sample(x)
- return x, short_cuts
- class ConvBlock(nn.Layer):
- def __init__(self, ch_in, ch_out):
- super(ConvBlock, self).__init__()
- self.conv = nn.Sequential(
- nn.Conv2D(
- ch_in, ch_out, kernel_size=3, stride=1, padding=1),
- nn.BatchNorm2D(ch_out),
- nn.ReLU(),
- nn.Conv2D(
- ch_out, ch_out, kernel_size=3, stride=1, padding=1),
- nn.BatchNorm2D(ch_out),
- nn.ReLU())
- def forward(self, x):
- return self.conv(x)
|