various fixes

aitlas/base/segmentation.py

@@ -1,12 +1,11 @@
 import logging
-import torch.nn as nn
 import torch.optim as optim
-import torch.nn.functional as nnf
 import torch
 
 from .models import BaseModel
 from .schemas import BaseSegmentationClassifierSchema
 from .metrics import SegmentationRunningScore
+from ..utils import DiceLoss
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 
@@ -25,8 +24,6 @@ def get_predicted(self, outputs, threshold=None):
         predicted = (predicted_probs >= (
             threshold if threshold else self.config.threshold
         )).long()
-        #predicted_probs = nnf.softmax(outputs, dim=1)
-        #predicted = (outputs >= predicted_probs).long()
         return predicted_probs, predicted
 
     def load_optimizer(self):
@@ -35,8 +32,7 @@ def load_optimizer(self):
 
     def load_criterion(self):
         """Load the loss function"""
-        return nn.BCEWithLogitsLoss()
-        #return nn.CrossEntropyLoss()
+        return DiceLoss()
 
     def load_lr_scheduler(self, optimizer):
         return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=5, factor=0.1, min_lr=1e-6)

aitlas/models/alexnet.py

@@ -34,9 +34,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class AlexNetMultiLabel(BaseMultilabelClassifier):
@@ -70,6 +70,6 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True

aitlas/models/convnext.py

@@ -28,9 +28,9 @@ def forward(self, x):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
     def extract_features(self):
         """ Remove final layers if we only need to extract features """
@@ -69,6 +69,6 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True

aitlas/models/densenet.py

@@ -44,9 +44,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class DenseNet161MultiLabel(BaseMultilabelClassifier):
@@ -88,6 +88,6 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True

aitlas/models/efficientnet.py

@@ -26,9 +26,9 @@ def forward(self, x):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
     """ Remove final layers if we only need to extract features """
     def extract_features(self):
@@ -65,9 +65,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class EfficientNetB4(BaseMulticlassClassifier):
@@ -92,9 +92,9 @@ def forward(self, x):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
     """ Remove final layers if we only need to extract features """
     def extract_features(self):
@@ -131,9 +131,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class EfficientNetB7(BaseMulticlassClassifier):
@@ -158,9 +158,9 @@ def forward(self, x):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.fc.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
     """ Remove final layers if we only need to extract features """
     def extract_features(self):
@@ -197,9 +197,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 

aitlas/models/resnet.py

@@ -30,7 +30,7 @@ def __init__(self, config):
                     weights=None, progress=False, num_classes=num_classes
                 )
                 # remove prefix "module."
-                checkpoint = {k.replace("module.backbone.", ""): v for k, v in checkpoint.items()}
+                checkpoint = {k.replace("backbone.", ""): v for k, v in checkpoint.items()}
                 checkpoint = {k.replace("module.", ""): v for k, v in checkpoint.items()}
                 for k, v in self.model.state_dict().items():
                     if k not in list(checkpoint):
@@ -56,9 +56,9 @@ def forward(self, x):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.fc.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
     def extract_features(self):
         """ Remove final layers if we only need to extract features """
@@ -95,9 +95,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.fc.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class ResNet50MultiLabel(BaseMultilabelClassifier):
@@ -139,9 +139,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.fc.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class ResNet152MultiLabel(BaseMultilabelClassifier):
@@ -172,6 +172,6 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.fc.parameters():
-            param.require_grad = True
+            param.requires_grad = True

aitlas/models/swin_transformer.py

@@ -18,6 +18,15 @@ def __init__(self, config):
             in_features=768, out_features=self.config.num_classes, bias=True
         )
 
+        if self.config.freeze:
+            self.freeze()
+
+    def freeze(self):
+        for param in self.model.parameters():
+            param.requires_grad = False
+        for param in self.model.head.parameters():
+            param.requires_grad = True
+
     def forward(self, x):
         return self.model(x)
 
@@ -36,5 +45,14 @@ def __init__(self, config):
             in_features=768, out_features=self.config.num_classes, bias=True
         )
 
+        if self.config.freeze:
+            self.freeze()
+
+    def freeze(self):
+        for param in self.model.parameters():
+            param.requires_grad = False
+        for param in self.model.head.parameters():
+            param.requires_grad = True
+
     def forward(self, x):
         return self.model(x)

aitlas/models/vgg.py

@@ -29,9 +29,9 @@ def forward(self, x):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
     def extract_features(self):
         """ Remove final layers if we only need to extract features """
@@ -70,9 +70,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class VGG16MultiLabel(BaseMultilabelClassifier):
@@ -105,9 +105,9 @@ def extract_features(self):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
 
 class VGG19MultiLabel(BaseMultilabelClassifier):
@@ -134,9 +134,9 @@ def forward(self, x):
 
     def freeze(self):
         for param in self.model.parameters():
-            param.require_grad = False
+            param.requires_grad = False
         for param in self.model.classifier.parameters():
-            param.require_grad = True
+            param.requires_grad = True
 
     def extract_features(self):
         """ Remove final layers if we only need to extract features """

aitlas/models/vision_transformer.py

@@ -31,7 +31,7 @@ def __init__(self, config):
                 in_features=768, out_features=num_classes, bias=True
             )
             # remove prefix "module."
-            checkpoint = {k.replace("module.backbone.", ""): v for k, v in checkpoint.items()}
+            checkpoint = {k.replace("backbone.", ""): v for k, v in checkpoint.items()}
             checkpoint = {k.replace("module.", ""): v for k, v in checkpoint.items()}
             for k, v in self.model.state_dict().items():
                 if k not in list(checkpoint):
@@ -51,6 +51,15 @@ def __init__(self, config):
                 in_features=768, out_features=self.config.num_classes, bias=True
             )
 
+        if self.config.freeze:
+            self.freeze()
+
+    def freeze(self):
+        for param in self.model.parameters():
+            param.requires_grad = False
+        for param in self.model.head.parameters():
+            param.requires_grad = True
+
     def forward(self, x):
         return self.model(x)
 
@@ -78,7 +87,7 @@ def __init__(self, config):
                 in_features=768, out_features=num_classes, bias=True
             )
             # remove prefix "module."
-            checkpoint = {k.replace("module.backbone.", ""): v for k, v in checkpoint.items()}
+            checkpoint = {k.replace("backbone.", ""): v for k, v in checkpoint.items()}
             checkpoint = {k.replace("module.", ""): v for k, v in checkpoint.items()}
             for k, v in self.model.state_dict().items():
                 if k not in list(checkpoint):
@@ -98,5 +107,14 @@ def __init__(self, config):
                 in_features=768, out_features=self.config.num_classes, bias=True
             )
 
+        if self.config.freeze:
+            self.freeze()
+
+    def freeze(self):
+        for param in self.model.parameters():
+            param.requires_grad = False
+        for param in self.model.head.parameters():
+            param.requires_grad = True
+
     def forward(self, x):
         return self.model(x)

aitlas/transforms/segmentation.py

@@ -24,6 +24,18 @@ def __call__(self, sample):
         return torch.tensor(sample, dtype=torch.float32) / 255
 
 
+class Pad(BaseTransforms):
+    def __call__(self, sample):
+        data_transforms = transforms.Compose(
+            [
+                transforms.ToPILImage(),
+                transforms.Pad(4),
+                transforms.ToTensor()
+            ]
+        )
+        return data_transforms(sample)
+
+
 class ColorTransformations(BaseTransforms):
     def __call__(self, sample):
         sample = np.asarray(sample)