formatted utils folder

Author: Johanmkr

utils/dataloaders/__init__.py

@@ -1,4 +1,4 @@
 __all__ = ["USPSDataset0_6", "USPSH5_Digit_7_9_Dataset"]
 
 from .usps_0_6 import USPSDataset0_6
-from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset
+from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset

utils/load_data.py

@@ -8,6 +8,6 @@ def load_data(dataset: str, *args, **kwargs) -> Dataset:
         case "usps_0-6":
             return USPSDataset0_6(*args, **kwargs)
         case "usps_7-9":
-            return  USPSH5_Digit_7_9_Dataset(*args, **kwargs)
+            return USPSH5_Digit_7_9_Dataset(*args, **kwargs)
         case _:
             raise ValueError(f"Dataset: {dataset} not implemented.")

utils/metrics/F1.py

@@ -84,4 +84,3 @@ def compute(self):
         )
 
         return f1_score
-

utils/metrics/__init__.py

@@ -3,3 +3,4 @@
 from .EntropyPred import EntropyPrediction
 from .F1 import F1Score
 from .recall import Recall
+from .precision import Precision

utils/metrics/precision.py

@@ -7,20 +7,23 @@
 
 
 class Precision(nn.Module):
-    """Metric module for precision. Can calculate precision both as a mean of precisions or as brute function of true positives and false positives. This is for now controller with the USE_MEAN macro.
+    """Metric module for precision. Can calculate precision both as a mean of precisions or as brute function of true positives and false positives.
 
     Parameters
     ----------
     num_classes : int
         Number of classes in the dataset.
+    use_mean : bool
+        Whether to calculate precision as a mean of precisions or as a brute function of true positives and false positives.
     """
 
-    def __init__(self, num_classes):
+    def __init__(self, num_classes: int, use_mean: bool = True):
         super().__init__()
 
         self.num_classes = num_classes
+        self.use_mean = use_mean
 
-    def forward(self, y_true, y_pred):
+    def forward(self, y_true: torch.tensor, y_pred: torch.tensor) -> torch.tensor:
         """Calculates the precision score given number of classes and the true and predicted labels.
 
         Parameters
@@ -43,7 +46,7 @@ def forward(self, y_true, y_pred):
             1, y_pred.unsqueeze(1), 1
         )
 
-        if USE_MEAN:
+        if self.use_mean:
             tp = torch.sum(true_oh * pred_oh, 0)
             fp = torch.sum(~true_oh.bool() * pred_oh, 0)
 
@@ -54,52 +57,5 @@ def forward(self, y_true, y_pred):
         return torch.nanmean(tp / (tp + fp))
 
 
-def test_precision_case1():
-    true_precision = 25.0 / 36 if USE_MEAN else 7.0 / 10
-
-    true1 = torch.tensor([0, 1, 2, 1, 0, 2, 1, 0, 2, 1])
-    pred1 = torch.tensor([0, 2, 1, 1, 0, 2, 0, 0, 2, 1])
-    P = Precision(3)
-    precision1 = P(true1, pred1)
-    assert precision1.allclose(torch.tensor(true_precision), atol=1e-5), (
-        f"Precision Score: {precision1.item()}"
-    )
-
-
-def test_precision_case2():
-    true_precision = 8.0 / 15 if USE_MEAN else 6.0 / 15
-
-    true2 = torch.tensor([0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4])
-    pred2 = torch.tensor([0, 0, 4, 3, 4, 0, 4, 4, 2, 3, 4, 1, 2, 4, 0])
-    P = Precision(5)
-    precision2 = P(true2, pred2)
-    assert precision2.allclose(torch.tensor(true_precision), atol=1e-5), (
-        f"Precision Score: {precision2.item()}"
-    )
-
-
-def test_precision_case3():
-    true_precision = 3.0 / 4 if USE_MEAN else 4.0 / 5
-
-    true3 = torch.tensor([0, 0, 0, 1, 0])
-    pred3 = torch.tensor([1, 0, 0, 1, 0])
-    P = Precision(2)
-    precision3 = P(true3, pred3)
-    assert precision3.allclose(torch.tensor(true_precision), atol=1e-5), (
-        f"Precision Score: {precision3.item()}"
-    )
-
-
-def test_for_zero_denominator():
-    true_precision = 0.0
-    true4 = torch.tensor([1, 1, 1, 1, 1])
-    pred4 = torch.tensor([0, 0, 0, 0, 0])
-    P = Precision(2)
-    precision4 = P(true4, pred4)
-    assert precision4.allclose(torch.tensor(true_precision), atol=1e-5), (
-        f"Precision Score: {precision4.item()}"
-    )
-
-
 if __name__ == "__main__":
     pass

utils/models/solveig_model.py

@@ -4,26 +4,26 @@
 
 class SolveigModel(nn.Module):
     """
-        A Convolutional Neural Network model for classification.
-
-         Args
-        ----
-        image_shape : tuple(int, int, int)
-            Shape of the input image (C, H, W).
-        num_classes : int
-            Number of classes in the dataset.
-
-        Attributes:
-        -----------
-        conv_block1 : nn.Sequential
-            First convolutional block containing a convolutional layer, ReLU activation, and max-pooling.
-        conv_block2 : nn.Sequential
-            Second convolutional block containing a convolutional layer and ReLU activation.
-        conv_block3 : nn.Sequential
-            Third convolutional block containing a convolutional layer and ReLU activation.
-        fc1 : nn.Linear
-            Fully connected layer that outputs the final classification scores.
-        """
+    A Convolutional Neural Network model for classification.
+
+     Args
+    ----
+    image_shape : tuple(int, int, int)
+        Shape of the input image (C, H, W).
+    num_classes : int
+        Number of classes in the dataset.
+
+    Attributes:
+    -----------
+    conv_block1 : nn.Sequential
+        First convolutional block containing a convolutional layer, ReLU activation, and max-pooling.
+    conv_block2 : nn.Sequential
+        Second convolutional block containing a convolutional layer and ReLU activation.
+    conv_block3 : nn.Sequential
+        Third convolutional block containing a convolutional layer and ReLU activation.
+    fc1 : nn.Linear
+        Fully connected layer that outputs the final classification scores.
+    """
 
     def __init__(self, image_shape, num_classes):
         super().__init__()
@@ -34,19 +34,19 @@ def __init__(self, image_shape, num_classes):
         self.conv_block1 = nn.Sequential(
             nn.Conv2d(in_channels=C, out_channels=25, kernel_size=3, padding=1),
             nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2)
+            nn.MaxPool2d(kernel_size=2, stride=2),
         )
 
         # Define the second convolutional block (conv + relu)
         self.conv_block2 = nn.Sequential(
             nn.Conv2d(in_channels=25, out_channels=50, kernel_size=3, padding=1),
-            nn.ReLU()
+            nn.ReLU(),
         )
 
         # Define the third convolutional block (conv + relu)
         self.conv_block3 = nn.Sequential(
             nn.Conv2d(in_channels=50, out_channels=100, kernel_size=3, padding=1),
-            nn.ReLU()
+            nn.ReLU(),
         )
 
         self.fc1 = nn.Linear(100 * 8 * 8, num_classes)
@@ -64,8 +64,7 @@ def forward(self, x):
 
 
 if __name__ == "__main__":
-
-    x = torch.randn(1,3, 16, 16)
+    x = torch.randn(1, 3, 16, 16)
 
     model = SolveigModel(x.shape[1:], 3)