attention_unet.py 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import paddle
  15. import paddle.nn as nn
  16. from paddlex.paddleseg.cvlibs import manager
  17. from paddlex.paddleseg.models import layers
  18. from paddlex.paddleseg import utils
  19. import numpy as np
  20. @manager.MODELS.add_component
  21. class AttentionUNet(nn.Layer):
  22. """
  23. The Attention-UNet implementation based on PaddlePaddle.
  24. As mentioned in the original paper, author proposes a novel attention gate (AG)
  25. that automatically learns to focus on target structures of varying shapes and sizes.
  26. Models trained with AGs implicitly learn to suppress irrelevant regions in an input image while
  27. highlighting salient features useful for a specific task.
  28. The original article refers to
  29. Oktay, O, et, al. "Attention u-net: Learning where to look for the pancreas."
  30. (https://arxiv.org/pdf/1804.03999.pdf).
  31. Args:
  32. num_classes (int): The unique number of target classes.
  33. pretrained (str, optional): The path or url of pretrained model. Default: None.
  34. """
  35. def __init__(self, num_classes, pretrained=None):
  36. super().__init__()
  37. n_channels = 3
  38. self.encoder = Encoder(n_channels, [64, 128, 256, 512])
  39. filters = np.array([64, 128, 256, 512, 1024])
  40. self.up5 = UpConv(ch_in=filters[4], ch_out=filters[3])
  41. self.att5 = AttentionBlock(
  42. F_g=filters[3], F_l=filters[3], F_out=filters[2])
  43. self.up_conv5 = ConvBlock(ch_in=filters[4], ch_out=filters[3])
  44. self.up4 = UpConv(ch_in=filters[3], ch_out=filters[2])
  45. self.att4 = AttentionBlock(
  46. F_g=filters[2], F_l=filters[2], F_out=filters[1])
  47. self.up_conv4 = ConvBlock(ch_in=filters[3], ch_out=filters[2])
  48. self.up3 = UpConv(ch_in=filters[2], ch_out=filters[1])
  49. self.att3 = AttentionBlock(
  50. F_g=filters[1], F_l=filters[1], F_out=filters[0])
  51. self.up_conv3 = ConvBlock(ch_in=filters[2], ch_out=filters[1])
  52. self.up2 = UpConv(ch_in=filters[1], ch_out=filters[0])
  53. self.att2 = AttentionBlock(
  54. F_g=filters[0], F_l=filters[0], F_out=filters[0] // 2)
  55. self.up_conv2 = ConvBlock(ch_in=filters[1], ch_out=filters[0])
  56. self.conv_1x1 = nn.Conv2D(
  57. filters[0], num_classes, kernel_size=1, stride=1, padding=0)
  58. self.pretrained = pretrained
  59. self.init_weight()
  60. def forward(self, x):
  61. x5, (x1, x2, x3, x4) = self.encoder(x)
  62. d5 = self.up5(x5)
  63. x4 = self.att5(g=d5, x=x4)
  64. d5 = paddle.concat([x4, d5], axis=1)
  65. d5 = self.up_conv5(d5)
  66. d4 = self.up4(d5)
  67. x3 = self.att4(g=d4, x=x3)
  68. d4 = paddle.concat((x3, d4), axis=1)
  69. d4 = self.up_conv4(d4)
  70. d3 = self.up3(d4)
  71. x2 = self.att3(g=d3, x=x2)
  72. d3 = paddle.concat((x2, d3), axis=1)
  73. d3 = self.up_conv3(d3)
  74. d2 = self.up2(d3)
  75. x1 = self.att2(g=d2, x=x1)
  76. d2 = paddle.concat((x1, d2), axis=1)
  77. d2 = self.up_conv2(d2)
  78. logit = self.conv_1x1(d2)
  79. logit_list = [logit]
  80. return logit_list
  81. def init_weight(self):
  82. if self.pretrained is not None:
  83. utils.load_entire_model(self, self.pretrained)
  84. class AttentionBlock(nn.Layer):
  85. def __init__(self, F_g, F_l, F_out):
  86. super().__init__()
  87. self.W_g = nn.Sequential(
  88. nn.Conv2D(
  89. F_g, F_out, kernel_size=1, stride=1, padding=0),
  90. nn.BatchNorm2D(F_out))
  91. self.W_x = nn.Sequential(
  92. nn.Conv2D(
  93. F_l, F_out, kernel_size=1, stride=1, padding=0),
  94. nn.BatchNorm2D(F_out))
  95. self.psi = nn.Sequential(
  96. nn.Conv2D(
  97. F_out, 1, kernel_size=1, stride=1, padding=0),
  98. nn.BatchNorm2D(1),
  99. nn.Sigmoid())
  100. self.relu = nn.ReLU()
  101. def forward(self, g, x):
  102. g1 = self.W_g(g)
  103. x1 = self.W_x(x)
  104. psi = self.relu(g1 + x1)
  105. psi = self.psi(psi)
  106. res = x * psi
  107. return res
  108. class UpConv(nn.Layer):
  109. def __init__(self, ch_in, ch_out):
  110. super().__init__()
  111. self.up = nn.Sequential(
  112. nn.Upsample(
  113. scale_factor=2, mode="bilinear"),
  114. nn.Conv2D(
  115. ch_in, ch_out, kernel_size=3, stride=1, padding=1),
  116. nn.BatchNorm2D(ch_out),
  117. nn.ReLU())
  118. def forward(self, x):
  119. return self.up(x)
  120. class Encoder(nn.Layer):
  121. def __init__(self, input_channels, filters):
  122. super().__init__()
  123. self.double_conv = nn.Sequential(
  124. layers.ConvBNReLU(input_channels, 64, 3),
  125. layers.ConvBNReLU(64, 64, 3))
  126. down_channels = filters
  127. self.down_sample_list = nn.LayerList([
  128. self.down_sampling(channel, channel * 2)
  129. for channel in down_channels
  130. ])
  131. def down_sampling(self, in_channels, out_channels):
  132. modules = []
  133. modules.append(nn.MaxPool2D(kernel_size=2, stride=2))
  134. modules.append(layers.ConvBNReLU(in_channels, out_channels, 3))
  135. modules.append(layers.ConvBNReLU(out_channels, out_channels, 3))
  136. return nn.Sequential(*modules)
  137. def forward(self, x):
  138. short_cuts = []
  139. x = self.double_conv(x)
  140. for down_sample in self.down_sample_list:
  141. short_cuts.append(x)
  142. x = down_sample(x)
  143. return x, short_cuts
  144. class ConvBlock(nn.Layer):
  145. def __init__(self, ch_in, ch_out):
  146. super(ConvBlock, self).__init__()
  147. self.conv = nn.Sequential(
  148. nn.Conv2D(
  149. ch_in, ch_out, kernel_size=3, stride=1, padding=1),
  150. nn.BatchNorm2D(ch_out),
  151. nn.ReLU(),
  152. nn.Conv2D(
  153. ch_out, ch_out, kernel_size=3, stride=1, padding=1),
  154. nn.BatchNorm2D(ch_out),
  155. nn.ReLU())
  156. def forward(self, x):
  157. return self.conv(x)