ruffed, isorted

Author: hzavadil98

utils/dataloaders/svhn.py

@@ -1,14 +1,14 @@
-from torch.utils.data import Dataset
+import os
+
 from scipy.io import loadmat
-import os 
+from torch.utils.data import Dataset
 from torchvision.datasets import SVHN
 
+
 class SVHNDataset(Dataset):
-    def __init__(self,
-                 datapath: str,
-                 transforms=None,
-                 download_data=True,
-                 split='train'):
+    def __init__(
+        self, datapath: str, transforms=None, download_data=True, split="train"
+    ):
         """
         Initializes the SVHNDataset object.
         Args:
@@ -20,36 +20,38 @@ def __init__(self,
             AssertionError: If the split is not 'train' or 'test'.
         """
         super().__init__()
-        assert split == 'train' or split == 'test'
-        
+        assert split == "train" or split == "test"
+
         if download_data:
             self._download_data(datapath, split)
-        
-        data = loadmat(os.path.join(datapath, f'{split}_32x32.mat'))
-        
+
+        data = loadmat(os.path.join(datapath, f"{split}_32x32.mat"))
+
         # Images on the form N x H x W x C
-        self.images = data['X'].transpose(3, 1, 0, 2)
-        self.labels = data['y'].flatten()
+        self.images = data["X"].transpose(3, 1, 0, 2)
+        self.labels = data["y"].flatten()
         self.labels[self.labels == 10] = 0
-        
+
         self.transforms = transforms
+
     def _download_data(self, path: str, split: str):
         """
         Downloads the SVHN dataset.
         Args:
             path (str): The directory where the dataset will be downloaded.
             split (str): The dataset split to download, either 'train' or 'test'.
         """
-        print(f'Downloading SVHN data into {path}')
-        SVHN(path, split=split, download=True) 
-    
+        print(f"Downloading SVHN data into {path}")
+        SVHN(path, split=split, download=True)
+
     def __len__(self):
         """
         Returns the number of samples in the dataset.
         Returns:
             int: The number of samples.
         """
         return len(self.labels)
+
     def __getitem__(self, index):
         """
         Retrieves the image and label at the specified index.
@@ -59,8 +61,8 @@ def __getitem__(self, index):
             tuple: A tuple containing the image and its corresponding label.
         """
         img, lab = self.images[index], self.labels[index]
-        
+
         if self.transforms is not None:
             img = self.transforms(img)
-        
-        return img, lab
+
+        return img, lab

utils/metrics/EntropyPred.py

@@ -4,21 +4,21 @@
 
 
 class EntropyPrediction(nn.Module):
-    def __init__(self, averages: str = 'average'):
+    def __init__(self, averages: str = "average"):
         """
         Initializes the EntropyPrediction module.
         Args:
-            averages (str): Specifies the method of aggregation for entropy values. 
+            averages (str): Specifies the method of aggregation for entropy values.
                             Must be either 'average' or 'sum'.
         Raises:
             AssertionError: If the averages parameter is not 'average' or 'sum'.
         """
         super().__init__()
-        
-        assert averages == 'average' or averages == 'sum'
+
+        assert averages == "average" or averages == "sum"
         self.averages = averages
         self.stored_entropy_values = []
-        
+
     def __call__(self, y_true, y_false_logits):
         """
         Computes the entropy between true labels and predicted logits, storing the results.