Update to handle nr_channels arg

Author: c-salomonsen

CollaborativeCoding/dataloaders/usps_0_6.py

@@ -8,6 +8,7 @@
 
 import h5py as h5
 import numpy as np
+import torch
 from PIL import Image
 from torch.utils.data import Dataset
 
@@ -83,6 +84,7 @@ def __init__(
         sample_ids: list,
         train: bool = False,
         transform=None,
+        nr_channels=1,
     ):
         super().__init__()
 
@@ -91,6 +93,7 @@ def __init__(
         self.transform = transform
         self.mode = "train" if train else "test"
         self.sample_ids = sample_ids
+        self.nr_channels = nr_channels
 
     def __len__(self):
         return len(self.sample_ids)
@@ -100,11 +103,18 @@ def __getitem__(self, id):
 
         with h5.File(self.filepath, "r") as f:
             data = f[self.mode]["data"][index].astype(np.uint8)
-            label = f[self.mode]["target"][index]
+            label = int(f[self.mode]["target"][index])
 
-        data = Image.fromarray(data, mode="L")
+        if self.nr_channels == 1:
+            data = Image.fromarray(data, mode="L")
+        elif self.nr_channels == 3:
+            data = Image.fromarray(data, mode="RGB")
+        else:
+            raise ValueError("Invalid number of channels")
 
         if self.transform:
             data = self.transform(data)
 
+        # label = torch.tensor(label).long()
+
         return data, label

CollaborativeCoding/load_data.py

@@ -89,23 +89,23 @@ def load_data(dataset: str, *args, **kwargs) -> tuple:
         sample_ids=train_samples,
         train=True,
         transform=transform,
-        nr_channels=kwargs.get("nr_channels"),
+        nr_channels=kwargs.get("nr_channels", 1),
     )
 
     val = dataset(
         data_path=data_dir,
         sample_ids=val_samples,
         train=True,
         transform=transform,
-        nr_channels=kwargs.get("nr_channels"),
+        nr_channels=kwargs.get("nr_channels", 1),
     )
 
     test = dataset(
         data_path=data_dir,
         sample_ids=test_samples,
         train=False,
         transform=transform,
-        nr_channels=kwargs.get("nr_channels"),
+        nr_channels=kwargs.get("nr_channels", 1),
     )
 
     return train, val, test