Add test for empty probability map

fepegar · Dec 22, 2020 · fabdec6 · fabdec6
1 parent aef479f
commit fabdec6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/tests/data/sampler/test_label_sampler.py b/tests/data/sampler/test_label_sampler.py
@@ -64,3 +64,14 @@ def test_no_labelmap(self):
         sampler = tio.LabelSampler(1)
         with self.assertRaises(RuntimeError):
             next(sampler(subject))
+
+    def test_empty_map(self):
+        # https://github.com/fepegar/torchio/issues/392
+        im = tio.ScalarImage(tensor=torch.rand(1, 6, 6, 6))
+        label = torch.zeros(1, 6, 6, 6)
+        label[..., 0] = 1  # voxels far from center
+        label_im = tio.LabelMap(tensor=label)
+        subject = tio.Subject(image=im, label=label_im)
+        sampler = tio.LabelSampler(4)
+        with self.assertRaises(RuntimeError):
+            next(sampler(subject))
diff --git a/tests/data/sampler/test_weighted_sampler.py b/tests/data/sampler/test_weighted_sampler.py
@@ -24,7 +24,7 @@ def get_sample(self, image_shape):
         subject = torchio.SubjectsDataset([subject])[0]
         return subject
 
-    def test_incosistent_shape(self):
+    def test_inconsistent_shape(self):
         # https://github.com/fepegar/torchio/issues/234#issuecomment-675029767
         subject = torchio.Subject(
             im1=torchio.ScalarImage(tensor=torch.rand(1, 4, 5, 6)),