net.py

@@ -7,7 +7,6 @@ from timm.models.layers import DropPath, to_2tuple, trunc_normal_
 from einops import rearrange
 
 from componets.SCSA import SCSA
-from componets.WTConvCV2 import WTConv2d
 
 
 def drop_path(x, drop_prob: float = 0., training: bool = False):
@@ -249,10 +248,6 @@ class DetailNode(nn.Module):
             self.theta_phi = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
             self.theta_rho = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
             self.theta_eta = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
-        elif useBlock == 2:
-            self.theta_phi = WTConv2d(in_channels=32, out_channels=32)
-            self.theta_rho = WTConv2d(in_channels=32, out_channels=32)
-            self.theta_eta = WTConv2d(in_channels=32, out_channels=32)
         else:
             self.theta_phi = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
             self.theta_rho = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
@@ -274,7 +269,7 @@ class DetailNode(nn.Module):
 class DetailFeatureFusion(nn.Module):
     def __init__(self, num_layers=3):
         super(DetailFeatureFusion, self).__init__()
-        INNmodules = [DetailNode(useBlock=2) for _ in range(num_layers)]
+        INNmodules = [DetailNode() for _ in range(num_layers)]
         self.net = nn.Sequential(*INNmodules)
 
     def forward(self, x):
@@ -286,7 +281,7 @@ class DetailFeatureFusion(nn.Module):
 class DetailFeatureExtraction(nn.Module):
     def __init__(self, num_layers=3):
         super(DetailFeatureExtraction, self).__init__()
-        INNmodules = [DetailNode(useBlock=2) for _ in range(num_layers)]
+        INNmodules = [DetailNode() for _ in range(num_layers)]
         self.net = nn.Sequential(*INNmodules)
 
     def forward(self, x):