|
1
|
|
|
import torch |
|
2
|
|
|
import torchio as tio |
|
3
|
|
|
from ...utils import TorchioTestCase |
|
4
|
|
|
|
|
5
|
|
|
|
|
6
|
|
|
class TestLabelSampler(TorchioTestCase): |
|
7
|
|
|
"""Tests for `LabelSampler` class.""" |
|
8
|
|
|
|
|
9
|
|
|
def test_label_sampler(self): |
|
10
|
|
|
sampler = tio.LabelSampler(5) |
|
11
|
|
|
for patch in sampler(self.sample_subject, num_patches=10): |
|
12
|
|
|
patch_center = patch['label'][tio.DATA][0, 2, 2, 2] |
|
13
|
|
|
self.assertEqual(patch_center, 1) |
|
14
|
|
|
|
|
15
|
|
|
def test_label_probabilities(self): |
|
16
|
|
|
labels = torch.Tensor((0, 0, 1, 1, 2, 1, 0)).reshape(1, 1, 1, -1) |
|
17
|
|
|
subject = tio.Subject( |
|
18
|
|
|
label=tio.Image(tensor=labels, type=tio.LABEL), |
|
19
|
|
|
) |
|
20
|
|
|
subject = tio.SubjectsDataset([subject])[0] |
|
21
|
|
|
probs_dict = {0: 0, 1: 50, 2: 25, 3: 25} |
|
22
|
|
|
sampler = tio.LabelSampler(5, 'label', label_probabilities=probs_dict) |
|
23
|
|
|
probabilities = sampler.get_probability_map(subject) |
|
24
|
|
|
fixture = torch.Tensor((0, 0, 2 / 12, 2 / 12, 3 / 12, 2 / 12, 0)) |
|
25
|
|
|
assert torch.all(probabilities.squeeze().eq(fixture)) |
|
26
|
|
|
|
|
27
|
|
|
def test_inconsistent_shape(self): |
|
28
|
|
|
# https://github.com/fepegar/torchio/issues/234#issuecomment-675029767 |
|
29
|
|
|
subject = tio.Subject( |
|
30
|
|
|
im1=tio.ScalarImage(tensor=torch.rand(2, 4, 5, 6)), |
|
31
|
|
|
im2=tio.LabelMap(tensor=torch.rand(1, 4, 5, 6)), |
|
32
|
|
|
) |
|
33
|
|
|
patch_size = 2 |
|
34
|
|
|
sampler = tio.LabelSampler(patch_size, 'im2') |
|
35
|
|
|
next(sampler(subject)) |
|
36
|
|
|
|
|
37
|
|
|
def test_multichannel_label_sampler(self): |
|
38
|
|
|
subject = tio.Subject( |
|
39
|
|
|
label=tio.LabelMap( |
|
40
|
|
|
tensor=torch.tensor( |
|
41
|
|
|
[ |
|
42
|
|
|
[[[1, 1]]], |
|
43
|
|
|
[[[0, 1]]] |
|
44
|
|
|
] |
|
45
|
|
|
) |
|
46
|
|
|
) |
|
47
|
|
|
) |
|
48
|
|
|
patch_size = 1 |
|
49
|
|
|
sampler = tio.LabelSampler( |
|
50
|
|
|
patch_size, |
|
51
|
|
|
'label', |
|
52
|
|
|
label_probabilities={0: 1, 1: 1} |
|
53
|
|
|
) |
|
54
|
|
|
# There are 2 voxels in the image, channels have same probabilities, |
|
55
|
|
|
# 1st voxel has probability 0.5 * 0.5 + 0 * 0.5 of being chosen while |
|
56
|
|
|
# 2nd voxel has probability 0.5 * 0.5 + 1 * 0.5 of being chosen. |
|
57
|
|
|
probabilities = sampler.get_probability_map(subject) |
|
58
|
|
|
fixture = torch.Tensor((1 / 4, 3 / 4)) |
|
59
|
|
|
assert torch.all(probabilities.squeeze().eq(fixture)) |
|
60
|
|
|
|
|
61
|
|
|
def test_no_labelmap(self): |
|
62
|
|
|
im = tio.ScalarImage(tensor=torch.rand(1, 1, 1, 1)) |
|
63
|
|
|
subject = tio.Subject(image=im, no_label=im) |
|
64
|
|
|
sampler = tio.LabelSampler(1) |
|
65
|
|
|
with self.assertRaises(RuntimeError): |
|
66
|
|
|
next(sampler(subject)) |
|
67
|
|
|
|
|
68
|
|
|
def test_empty_map(self): |
|
69
|
|
|
# https://github.com/fepegar/torchio/issues/392 |
|
70
|
|
|
im = tio.ScalarImage(tensor=torch.rand(1, 6, 6, 6)) |
|
71
|
|
|
label = torch.zeros(1, 6, 6, 6) |
|
72
|
|
|
label[..., 0] = 1 # voxels far from center |
|
73
|
|
|
label_im = tio.LabelMap(tensor=label) |
|
74
|
|
|
subject = tio.Subject(image=im, label=label_im) |
|
75
|
|
|
sampler = tio.LabelSampler(4) |
|
76
|
|
|
with self.assertRaises(RuntimeError): |
|
77
|
|
|
next(sampler(subject)) |
|
78
|
|
|
|
fepegar /
torchio
