Formatting

Author: c-salomonsen

utils/dataloaders/__init__.py

@@ -1,6 +1,11 @@
-__all__ = ["USPSDataset0_6", "USPSH5_Digit_7_9_Dataset", "MNISTDataset0_3", "SVHNDataset"]
+__all__ = [
+    "USPSDataset0_6",
+    "USPSH5_Digit_7_9_Dataset",
+    "MNISTDataset0_3",
+    "SVHNDataset",
+]
 
 from .mnist_0_3 import MNISTDataset0_3
+from .svhn import SVHNDataset
 from .usps_0_6 import USPSDataset0_6
 from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset
-from .svhn import SVHNDataset

utils/dataloaders/svhn.py

@@ -1,4 +1,5 @@
 import os
+
 import numpy as np
 from scipy.io import loadmat
 from torch.utils.data import Dataset
@@ -7,13 +8,13 @@
 
 class SVHNDataset(Dataset):
     def __init__(
-        self, 
-        data_path: str, 
+        self,
+        data_path: str,
         train: bool,
-        transform=None, 
-        download:bool=True, 
-        nr_channels=3
-        ):
+        transform=None,
+        download: bool = True,
+        nr_channels=3,
+    ):
         """
         Initializes the SVHNDataset object.
         Args:
@@ -26,8 +27,8 @@ def __init__(
         """
         super().__init__()
         # assert split == "train" or split == "test"
-        self.split = 'train' if train else 'test'
-    
+        self.split = "train" if train else "test"
+
         if download:
             self._download_data(data_path)
 
@@ -37,7 +38,7 @@ def __init__(
         self.images = data["X"].transpose(3, 1, 0, 2)
         self.labels = data["y"].flatten()
         self.labels[self.labels == 10] = 0
-        
+
         self.nr_channels = nr_channels
         self.transforms = transform
 
@@ -49,7 +50,7 @@ def _download_data(self, path: str):
             split (str): The dataset split to download, either 'train' or 'test'.
         """
         print(f"Downloading SVHN data into {path}")
-        
+
         SVHN(path, split=self.split, download=True)
 
     def __len__(self):
@@ -72,7 +73,7 @@ def __getitem__(self, index):
 
         if self.nr_channels == 1:
             img = np.mean(img, axis=2, keepdims=True)
-        
+
         if self.transforms is not None:
             img = self.transforms(img)
 

utils/load_data.py

@@ -1,6 +1,11 @@
 from torch.utils.data import Dataset
 
-from .dataloaders import MNISTDataset0_3, USPSDataset0_6, USPSH5_Digit_7_9_Dataset, SVHNDataset
+from .dataloaders import (
+    MNISTDataset0_3,
+    SVHNDataset,
+    USPSDataset0_6,
+    USPSH5_Digit_7_9_Dataset,
+)
 
 
 def load_data(dataset: str, *args, **kwargs) -> Dataset: