adjusted my dataloader to the new format, added test for my model

Author: sot176

tests/test_models.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from utils.models import ChristianModel, JanModel
+from utils.models import ChristianModel, JanModel, SolveigModel
 
 
 @pytest.mark.parametrize(
@@ -33,3 +33,18 @@ def test_jan_model(image_shape, num_classes):
 
     assert y.shape == (n, num_classes), f"Shape: {y.shape}"
 
+
+@pytest.mark.parametrize(
+    "image_shape, num_classes",
+    [((3, 16, 16), 3), ((3, 16, 16), 7)],
+)
+def test_solveig_model(image_shape, num_classes):
+    n, c, h, w = 5, *image_shape
+
+    model = SolveigModel(image_shape, num_classes)
+
+    x = torch.randn(n, c, h, w)
+    y = model(x)
+
+    assert y.shape == (n, num_classes), f"Shape: {y.shape}"
+

utils/dataloaders/uspsh5_7_9.py

@@ -4,6 +4,7 @@
 from PIL import Image
 from torch.utils.data import Dataset
 from torchvision import transforms
+from pathlib import Path
 
 
 class USPSH5_Digit_7_9_Dataset(Dataset):
@@ -30,7 +31,7 @@ class USPSH5_Digit_7_9_Dataset(Dataset):
         A transform function to apply to the images.
     """
 
-    def __init__(self, h5_path, mode, transform=None):
+    def __init__(self, data_path, train = False, transform=None):
         super().__init__()
         """
         Initializes the USPS dataset by loading images and labels from the given `.h5` file.
@@ -43,12 +44,13 @@ def __init__(self, h5_path, mode, transform=None):
         transform : callable, optional, default=None
             A transform function to apply on images.
         """
-
+        self.filename = "usps.h5"
+        path = data_path if isinstance(data_path, Path) else Path(data_path)
+        self.filepath = path / self.filename
         self.transform = transform
-        self.mode = mode
-        self.h5_path = h5_path
+        self.mode = "train" if train else "test"
         # Load the dataset from the HDF5 file
-        with h5py.File(self.h5_path, "r") as hf:
+        with h5py.File(self.filepath, "r") as hf:
             images = hf[self.mode]["data"][:]
             labels = hf[self.mode]["target"][:]
 
@@ -105,8 +107,8 @@ def main():
 
     # Load the dataset
     dataset = USPSH5_Digit_7_9_Dataset(
-        h5_path="C:/Users/Solveig/OneDrive/Dokumente/UiT PhD/Courses/Git/usps.h5",
-        mode="train",
+        data_path="C:/Users/Solveig/OneDrive/Dokumente/UiT PhD/Courses/Git",
+        train = False,
         transform=transform,
     )
     data_loader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True)