add stanet

Author: LinoSun

change_detection_pytorch/stanet/BAM.py

@@ -10,11 +10,11 @@ class BAM(nn.Module):
     def __init__(self, in_dim, ds=8, activation=nn.ReLU):
         super(BAM, self).__init__()
         self.chanel_in = in_dim
-        self.key_channel = self.chanel_in //8
+        self.key_channel = self.chanel_in // 8
         self.activation = activation
         self.ds = ds  #
         self.pool = nn.AvgPool2d(self.ds)
-        print('ds: ',ds)
+        print('ds: ', ds)
         self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
         self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
         self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
@@ -35,7 +35,7 @@ def forward(self, input):
         proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1)  # B X C X (N)/(ds*ds)
         proj_key = self.key_conv(x).view(m_batchsize, -1, width * height)  # B X C x (*W*H)/(ds*ds)
         energy = torch.bmm(proj_query, proj_key)  # transpose check
-        energy = (self.key_channel**-.5) * energy
+        energy = (self.key_channel ** -.5) * energy
 
         attention = self.softmax(energy)  # BX (N) X (N)/(ds*ds)/(ds*ds)
 
@@ -44,9 +44,7 @@ def forward(self, input):
         out = torch.bmm(proj_value, attention.permute(0, 2, 1))
         out = out.view(m_batchsize, C, width, height)
 
-        out = F.interpolate(out, [width*self.ds,height*self.ds])
+        out = F.interpolate(out, [width * self.ds, height * self.ds])
         out = out + input
 
         return out
-
-

change_detection_pytorch/stanet/model.py

@@ -27,6 +27,7 @@ class STANet(torch.nn.Module):
         https://www.mdpi.com/2072-4292/12/10/1662
 
     """
+
     def __init__(
             self,
             encoder_name: str = "resnet",