fix!: new model

models.py

@@ -135,6 +135,76 @@ def ResNet101(num_classes=1000):
 def ResNet152(num_classes=1000):
     return ResNet(Bottleneck, [3, 8, 36, 3], num_classes)
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SAM(nn.Module):
+    def __init__(self, bias=False):
+        super(SAM, self).__init__()
+        self.bias = bias
+        self.conv = nn.Conv2d(
+            in_channels=2,
+            out_channels=1,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+            dilation=1,
+            bias=self.bias,
+        )
+
+    def forward(self, x):
+        max = torch.max(x, 1)[0].unsqueeze(1)
+        avg = torch.mean(x, 1).unsqueeze(1)
+        concat = torch.cat((max, avg), dim=1)
+        output = self.conv(concat)
+        output = F.sigmoid(output) * x
+        return output
+
+
+class CAM(nn.Module):
+    def __init__(self, channels, r):
+        super(CAM, self).__init__()
+        self.channels = channels
+        self.r = r
+        self.linear = nn.Sequential(
+            nn.Linear(
+                in_features=self.channels,
+                out_features=self.channels // self.r,
+                bias=True,
+            ),
+            nn.ReLU(inplace=True),
+            nn.Linear(
+                in_features=self.channels // self.r,
+                out_features=self.channels,
+                bias=True,
+            ),
+        )
+
+    def forward(self, x):
+        max = F.adaptive_max_pool2d(x, output_size=1)
+        avg = F.adaptive_avg_pool2d(x, output_size=1)
+        b, c, _, _ = x.size()
+        linear_max = self.linear(max.view(b, c)).view(b, c, 1, 1)
+        linear_avg = self.linear(avg.view(b, c)).view(b, c, 1, 1)
+        output = linear_max + linear_avg
+        output = F.sigmoid(output) * x
+        return output
+
+
+class CBAM(nn.Module):
+    def __init__(self, channels, r):
+        super(CBAM, self).__init__()
+        self.channels = channels
+        self.r = r
+        self.sam = SAM(bias=False)
+        self.cam = CAM(channels=self.channels, r=self.r)
+
+    def forward(self, x):
+        output = self.cam(x)
+        output = self.sam(output)
+        return output + x
 
 class ClassifierHead(nn.Module):
     def __init__(self, in_features, num_classes):
@@ -165,6 +235,9 @@ class ResNetUNet(ResNet):
     def __init__(self, block, num_blocks, num_classes=1000):
         super().__init__(block, num_blocks, num_classes)
 
+        # Get the expansion factor
+        expansion = block.expansion
+
         # Calculate encoder channel sizes
         self.enc_channels = [
             64,
@@ -174,23 +247,14 @@ class ResNetUNet(ResNet):
             512 * block.expansion,
         ]
 
-        # Replace t_max_avg_pooling with standard avgpool
         in_features = 512 * block.expansion
         self.classifier_head = ClassifierHead(in_features, num_classes)
 
-        # Decoder layers remain the same
-        self.decoder5 = nn.Sequential(
-            nn.Conv2d(2048 + 1024, 1024, 3, padding=1),
-            nn.BatchNorm2d(1024),
-            nn.ReLU(inplace=True),
-            nn.Conv2d(1024, 512, 3, padding=1),
-            nn.BatchNorm2d(512),
-            nn.ReLU(inplace=True),
-            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
-        )
+        self.cbam = CBAM(channels=512 * block.expansion, r=16)
 
-        self.decoder4 = nn.Sequential(
-            nn.Conv2d(512 + 512, 512, 3, padding=1),
+        # Calculate encoder channel sizes
+        self.decoder5 = nn.Sequential(
+            nn.Conv2d((512 * expansion) + (256 * expansion), 512, 3, padding=1),
             nn.BatchNorm2d(512),
             nn.ReLU(inplace=True),
             nn.Conv2d(512, 256, 3, padding=1),
@@ -199,8 +263,8 @@ class ResNetUNet(ResNet):
             nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
         )
 
-        self.decoder3 = nn.Sequential(
-            nn.Conv2d(256 + 256, 256, 3, padding=1),
+        self.decoder4 = nn.Sequential(
+            nn.Conv2d(256 + (128 * expansion), 256, 3, padding=1),
             nn.BatchNorm2d(256),
             nn.ReLU(inplace=True),
             nn.Conv2d(256, 128, 3, padding=1),
@@ -209,8 +273,8 @@ class ResNetUNet(ResNet):
             nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
         )
 
-        self.decoder2 = nn.Sequential(
-            nn.Conv2d(128 + 64, 128, 3, padding=1),
+        self.decoder3 = nn.Sequential(
+            nn.Conv2d(128 + (64 * expansion), 128, 3, padding=1),
             nn.BatchNorm2d(128),
             nn.ReLU(inplace=True),
             nn.Conv2d(128, 64, 3, padding=1),
@@ -219,8 +283,18 @@ class ResNetUNet(ResNet):
             nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
         )
 
+        self.decoder2 = nn.Sequential(
+            nn.Conv2d(64 + 64, 64, 3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, 3, padding=1), 
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+        )
+
         self.final_conv = nn.Sequential(
-            nn.Conv2d(64, 32, 3, padding=1),
+            nn.Conv2d(64, 32, 3, padding=1),  
             nn.BatchNorm2d(32),
             nn.ReLU(inplace=True),
             nn.Conv2d(32, 1, 1),
@@ -265,6 +339,8 @@ class ResNetUNet(ResNet):
             seg_out, size=input_size, mode="bilinear", align_corners=True
         )
 
+        attended_features = self.cbam(e5)
+
         # Use segmentation to mask features before classification
         # Upsample segmentation mask to match feature size
         attention_mask = F.interpolate(
@@ -272,15 +348,13 @@ class ResNetUNet(ResNet):
         )
 
         # Apply attention mask to features
-        attended_features = e5 * (0.25 + attention_mask)
+        attended_features = attended_features * (0.25 + attention_mask)
 
-        # Use new classifier head
         cls_out = self.classifier_head(attended_features)
 
         return cls_out, seg_out
 
 
-# Helper functions without K and T parameters
 def ResNet18UNet(num_classes=1000):
     return ResNetUNet(BasicBlock, [2, 2, 2, 2], num_classes)
 
@@ -298,4 +372,4 @@ def ResNet101UNet(num_classes=1000):
 
 
 def ResNet152UNet(num_classes=1000):
-    return ResNetUNet(Bottleneck, [3, 8, 36, 3], num_classes)
+    return ResNetUNet(Bottleneck, [3, 8, 36, 3], num_classes)