Passed
Pull Request — master (#226)
by Fernando
01:12
created

tests.utils.TorchioTestCase.get_image_path()   A

Complexity

Conditions 3

Size

Total Lines 18
Code Lines 17

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 3
eloc 17
nop 5
dl 0
loc 18
rs 9.55
c 0
b 0
f 0
1
import copy
2
import shutil
3
import random
4
import tempfile
5
import unittest
6
from pathlib import Path
7
import torch
8
import numpy as np
9
import nibabel as nib
10
from torchio.datasets import IXITiny
11
from torchio import INTENSITY, LABEL, DATA, Image, ImagesDataset, Subject
12
13
14
class TorchioTestCase(unittest.TestCase):
15
16
    def setUp(self):
17
        """Set up test fixtures, if any."""
18
        self.dir = Path(tempfile.gettempdir()) / '.torchio_tests'
19
        self.dir.mkdir(exist_ok=True)
20
        random.seed(42)
21
        np.random.seed(42)
22
23
        registration_matrix = np.array([
24
            [1, 0, 0, 10],
25
            [0, 1, 0, 0],
26
            [0, 0, 1.2, 0],
27
            [0, 0, 0, 1]
28
        ])
29
30
        subject_a = Subject(
31
            t1=Image(self.get_image_path('t1_a'), INTENSITY),
32
        )
33
        subject_b = Subject(
34
            t1=Image(self.get_image_path('t1_b'), INTENSITY),
35
            label=Image(self.get_image_path('label_b', binary=True), LABEL),
36
        )
37
        subject_c = Subject(
38
            label=Image(self.get_image_path('label_c', binary=True), LABEL),
39
        )
40
        subject_d = Subject(
41
            t1=Image(
42
                self.get_image_path('t1_d'),
43
                INTENSITY,
44
                pre_affine=registration_matrix,
45
            ),
46
            t2=Image(self.get_image_path('t2_d'), INTENSITY),
47
            label=Image(self.get_image_path('label_d', binary=True), LABEL),
48
        )
49
        self.subjects_list = [
50
            subject_a,
51
            subject_b,
52
            subject_c,
53
            subject_d,
54
        ]
55
        self.dataset = ImagesDataset(self.subjects_list)
56
        self.sample = self.dataset[-1]
57
58
    def tearDown(self):
59
        """Tear down test fixtures, if any."""
60
        print('Deleting', self.dir)
61
        shutil.rmtree(self.dir)
62
63
    def assertTensorEqual(self, a, b):
64
        assert torch.all(torch.eq(a, b))
65
66
    def make_2d(self, sample):
67
        sample = copy.deepcopy(sample)
68
        for image in sample.get_images(intensity_only=False):
69
            image[DATA] = image[DATA][:, 0:1, ...]
0 ignored issues
show
Comprehensibility Best Practice introduced by
The variable DATA does not seem to be defined.
Loading history...
70
        return sample
71
72
    def get_inconsistent_sample(self):
73
        """Return a sample containing images of different shape."""
74
        subject = Subject(
75
            t1=Image(self.get_image_path('t1_inc'), INTENSITY),
76
            t2=Image(
77
                self.get_image_path('t2_inc', shape=(10, 20, 31)), INTENSITY),
78
            label=Image(
79
                self.get_image_path(
80
                    'label_inc',
81
                    shape=(8, 17, 25),
82
                    binary=True,
83
                ),
84
                LABEL,
85
            ),
86
        )
87
        subjects_list = [subject]
88
        dataset = ImagesDataset(subjects_list)
89
        return dataset[0]
90
91
    def get_reference_image_and_path(self):
92
        """Return a reference image and its path"""
93
        path = self.get_image_path('ref', shape=(10, 20, 31), spacing=(1, 1, 2))
94
        image = Image(path, INTENSITY)
95
        return image, path
96
97
    def get_ixi_tiny(self):
98
        root_dir = Path(tempfile.gettempdir()) / 'torchio' / 'ixi_tiny'
99
        return IXITiny(root_dir, download=True)
100
101
    def get_image_path(
102
            self,
103
            stem,
104
            binary=False,
105
            shape=(10, 20, 30),
106
            spacing=(1, 1, 1),
107
            ):
108
        data = np.random.rand(*shape)
109
        if binary:
110
            data = (data > 0.5).astype(np.uint8)
111
        affine = np.diag((*spacing, 1))
112
        suffix = random.choice(('.nii.gz', '.nii', '.nrrd', '.minc', '.img'))
113
        path = self.dir / f'{stem}{suffix}'
114
        if np.random.rand() > 0.5:
115
            path = str(path)
116
        image = Image(tensor=data, affine=affine)
117
        image.save(path)
118
        return path
119