Merge pull request #61 from SFI-Visual-Intelligence/mag-branch

Mondays sync

Author: Seilmast

utils/arg_parser.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.",
     )
 
@@ -69,6 +70,21 @@ def get_args():
         nargs="+",
         help="Which metric to use for evaluation",
     )
+    
+    parser.add_argument(
+        '--imagesize',
+        type=int,
+        default=28,
+        help='Imagesize'
+    )
+    
+    parser.add_argument(
+        '--nr_channels',
+        type=int,
+        default=1,
+        choices=[1,3],
+        help='Number of image channels'
+    )
 
     # Training specific values
     parser.add_argument(

utils/dataloaders/svhn.py

@@ -1,51 +1,56 @@
 import os
-
+import numpy as np
 from scipy.io import loadmat
 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"
+        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"
-
-        if download_data:
-            self._download_data(datapath, split)
+        # assert split == "train" or split == "test"
+        self.split = 'train' if train else 'test'
+    
+        if download:
+            self._download_data(data_path)
 
-        data = loadmat(os.path.join(datapath, f"{split}_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 = transform
 
-        self.transforms = transforms
-
-    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=split, download=True)
+        
+        SVHN(path, split=self.split, download=True)
 
     def __len__(self):
         """
@@ -65,6 +70,9 @@ def __getitem__(self, index):
         """
         img, lab = self.images[index], self.labels[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

@@ -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/models/magnus_model.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()
                     ]))