a b/test/test_utils.py
1
from unittest import TestCase
2
3
import nibabel as nib
4
import numpy as np
5
6
from fetal_net.utils.utils import resize
7
from fetal_net.utils.sitk_utils import resample_to_spacing
8
9
10
class TestUtils(TestCase):
11
    def _resize_image_test(self, image, target_shape):
12
        original_image_shape = image.shape
13
        new_image = resize(image, target_shape)
14
        self.assertEqual(new_image.shape, target_shape)
15
        new_image = resize(new_image, original_image_shape, interpolation="linear")
16
        self.assertEqual(new_image.shape, original_image_shape)
17
18
    def _create_image(self, image_shape):
19
        data = np.asarray(np.arange(np.prod(image_shape)).reshape(image_shape), dtype=np.float)
20
        affine = np.zeros((4, 4))
21
        np.fill_diagonal(affine, 1)
22
        return nib.Nifti1Image(data, affine)
23
24
    def test_resize_image_1(self):
25
        image_shape = (4, 4, 4)
26
        image = self._create_image(image_shape)
27
        new_size = (2, 2, 2)
28
        self._resize_image_test(image, new_size)
29
30
    def test_resize_image_2(self):
31
        self._resize_image_test(self._create_image((12, 10, 8)), (8, 8, 8))
32
33
    def test_resize_image_2d(self):
34
        data = np.arange(1, 5).reshape((2, 2))
35
        new_data = resample_to_spacing(data, (2, 2), (1, 1), interpolation="nearest")
36
        self.assertTrue(np.all(new_data == np.asarray([[1, 1, 2, 2],
37
                                                       [1, 1, 2, 2],
38
                                                       [3, 3, 4, 4],
39
                                                       [3, 3, 4, 4]])))
40
        orig_data = resample_to_spacing(new_data, (1, 1), (2, 2), interpolation="linear")
41
        self.assertTrue(np.all(data == orig_data))
42
43
    def test_resize_image_3(self):
44
        self._resize_image_test(self._create_image((2, 5, 3)), (7, 5, 11))
45
46
    def test_resize_image_3d(self):
47
        data = np.arange(1, 9).reshape((2, 2, 2))
48
        new_data = resample_to_spacing(data, (2, 2, 2), (1, 1, 1), interpolation="nearest")
49
        self.assertTrue(np.all(new_data[0] == np.asarray([[1, 1, 2, 2],
50
                                                          [1, 1, 2, 2],
51
                                                          [3, 3, 4, 4],
52
                                                          [3, 3, 4, 4]])))
53
        orig_data = resample_to_spacing(new_data, (1, 1, 1), (2, 2, 2), interpolation="linear")
54
        self.assertTrue(np.all(data == orig_data))
55
56
    def test_images_align(self):
57
        data = np.arange(1, 9).reshape((2, 2, 2))
58
        affine = np.diag(np.ones(4) * 2)
59
        affine[3, 3] = 1
60
        image_nib = nib.Nifti1Image(data, affine=affine)
61
        new_image_nib = resize(image_nib, (4, 4, 4), interpolation="nearest")
62
        self.assertTrue(np.all(new_image_nib.get_data()[0] == np.asarray([[1, 1, 2, 2],
63
                                                                          [1, 1, 2, 2],
64
                                                                          [3, 3, 4, 4],
65
                                                                          [3, 3, 4, 4]])))
66
        self.assertTrue(np.all(new_image_nib.affine == np.asarray([[1., 0., 0., -0.5],
67
                                                                   [0., 1., 0., -0.5],
68
                                                                   [0., 0., 1., -0.5],
69
                                                                   [0., 0., 0., 1.]])))