feat: ✨ Dataloading works. Ready to train.

janelia-cellmap · Apr 2, 2024 · d83d1f3 · d83d1f3
1 parent 85fd727
commit d83d1f3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/src/cellmap_data/dataset.py b/src/cellmap_data/dataset.py
@@ -138,9 +138,7 @@ def __getitem__(self, idx):
         outputs = {}
         for array_name in self.input_arrays.keys():
             self.input_sources[array_name].set_spatial_transforms(spatial_transforms)
-            outputs[array_name] = self.input_sources[array_name][center][
-                None, None, ...
-            ]
+            outputs[array_name] = self.input_sources[array_name][center][None, ...]
         # TODO: Allow for distribtion of array gathering to multiple threads
         for array_name in self.target_arrays.keys():
             class_arrays = []
@@ -149,7 +147,7 @@ def __getitem__(self, idx):
                     spatial_transforms
                 )
                 class_arrays.append(self.target_sources[array_name][label][center])
-            outputs[array_name] = torch.stack(class_arrays)[None, ...]
+            outputs[array_name] = torch.stack(class_arrays)
         return outputs
 
     def __iter__(self):

diff --git a/src/cellmap_data/multidataset.py b/src/cellmap_data/multidataset.py
@@ -25,6 +25,7 @@ def __init__(
         target_arrays: dict[str, dict[str, Sequence[int | float]]],
         datasets: Iterable[CellMapDataset],
     ):
+        super().__init__(datasets)
         self.input_arrays = input_arrays
         self.target_arrays = target_arrays
         self.classes = classes