update workflows

nitrain · May 6, 2024 · 4014b02 · 4014b02
1 parent 1fefb9e
commit 4014b02
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/nitrain/transforms/labels.py b/nitrain/transforms/labels.py
@@ -17,7 +17,7 @@ class ExpandLabels(BaseTransform):
     It is also possible to keep the original values in the channels 
     instead of making all values equal to 1.
     """
-    def __init__(self, keep_values=False, as_channels=False):
+    def __init__(self, keep_values=False, as_channels=True):
         self.keep_values = keep_values
         self.as_channels = as_channels
 

diff --git a/tests/test_workflows.py b/tests/test_workflows.py
@@ -78,12 +78,12 @@ def test_multiclass_segmentation(self):
 
         x,y = dataset[0]
 
-       # data_train, data_test = dataset.split(0.8)
+        data_train, data_test = dataset.split(0.8)
 
-        #loader = nt.Loader(dataset,
-        #                   images_per_batch=4)
-        #
-        #xb, yb = next(iter(loader))
+        loader = nt.Loader(data_train,
+                           images_per_batch=4)
+
+        xb, yb = next(iter(loader))
 #
         #arch_fn = nt.fetch_architecture('unet', dim=2)
         #model = arch_fn(x.shape[:-1]+(1,),