...

Author: Seilmast

test.py

@@ -35,7 +35,8 @@ def get_args():
 
     parser.add_argument(
         "--download-data",
-        action="store_true",
+        type=bool,
+        default=False,
         help="Whether the data should be downloaded or not. Might cause code to start a bit slowly.",
     )
 

utils/arg_parser.py

@@ -7,46 +7,50 @@
 
 class SVHNDataset(Dataset):
     def __init__(
-        self, datapath: str, 
-        transforms=None, 
-        download_data=True, 
+        self, 
+        data_path: str, 
+        train: bool,
+        transform=None, 
+        download:bool=True, 
         nr_channels=3
         ):
         """
         Initializes the SVHNDataset object.
         Args:
-            datapath (str): Path to where the data lies. If download_data is set to True, this is where the data will be downloaded.
+            data_path (str): Path to where the data lies. If download_data is set to True, this is where the data will be downloaded.
             transforms: Torch composite of transformations which are to be applied to the dataset images.
-            download_data (bool): If True, downloads the dataset to the specified datapath.
+            download_data (bool): If True, downloads the dataset to the specified data_path.
             split (str): The dataset split to use, either 'train' or 'test'.
         Raises:
             AssertionError: If the split is not 'train' or 'test'.
         """
         super().__init__()
         # assert split == "train" or split == "test"
+        self.split = 'train' if train else 'test'
+    
+        if download:
+            self._download_data(data_path)
 
-        if download_data:
-            self._download_data(datapath)
-
-        data = loadmat(os.path.join(datapath, f"train_32x32.mat"))
+        data = loadmat(os.path.join(data_path, f"{self.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.labels[self.labels == 10] = 0
 
         self.nr_channels = nr_channels
-        self.transforms = transforms
+        self.transforms = transform
 
-    def _download_data(self, path: str, split: str):
+    def _download_data(self, path: 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='train', download=True)
+        
+        SVHN(path, split=self.split, download=True)
 
     def __len__(self):
         """

utils/dataloaders/svhn.py

@@ -41,7 +41,7 @@ def load_data(dataset: str, *args, **kwargs) -> Dataset:
         case "usps_7-9":
             return USPSH5_Digit_7_9_Dataset(*args, **kwargs)
         case "svhn":
-            raise SVHNDataset(*args, **kwargs)
+            return SVHNDataset(*args, **kwargs)
         case "mnist_4-9":
             raise NotImplementedError("MNIST 4-9 dataset not yet implemented.")
         case _:

utils/load_data.py

@@ -2,24 +2,24 @@
 
 
 class MagnusModel(nn.Module):
-    def __init__(self, imagesize: int, imagechannels: int, n_classes: int = 10):
+    def __init__(self, image_shape: int, num_classes: int, imagechannels: int):
         """
         Magnus model contains the model for Magnus' part of the homeexam.
         This class contains a neural network consisting of three linear layers of 133 neurons each,
         with ReLU activation between each layer.
 
         Args
         ----
-            imagesize (int): Expected size of input image. This is needed to scale first layer input
+            image_shape (int): Expected size of input image. This is needed to scale first layer input
             imagechannels (int): Expected number of image channels. This is needed to scale first layer input
-            n_classes (int): Number of classes we are to provide.
+            num_classes (int): Number of classes we are to provide.
 
         Returns
         -------
             MagnusModel (nn.Module): Neural network as described above in this docstring.
         """
         super().__init__()
-        self.imagesize = imagesize
+        self.image_shape = image_shape
         self.imagechannels = imagechannels
 
         self.layer1 = nn.Sequential(*([
@@ -31,7 +31,7 @@ def __init__(self, imagesize: int, imagechannels: int, n_classes: int = 10):
                         nn.ReLU()
                     ]))
         self.layer3 = nn.Sequential(*([
-                        nn.Linear(133, n_classes), 
+                        nn.Linear(133, num_classes), 
                         nn.ReLU()
                     ]))