tests.utils

Complexity

Total Complexity

39

Size/Duplication

Total Lines	259
Duplicated Lines	0 %

Importance

Changes

0

1 Function

Rating	Name	Duplication	Size	Complexity
A	get_all_random_transforms()	0	6	1

import copy
import os
import random
import shutil
import tempfile
import unittest
from collections.abc import Sequence
from pathlib import Path
from random import shuffle

import numpy as np
import pytest
import torch

import torchio as tio


class TorchioTestCase(unittest.TestCase):
    def setUp(self):
        """Set up test fixtures, if any."""
        self.dir = Path(tempfile.gettempdir()) / os.urandom(24).hex()
        self.dir.mkdir(exist_ok=True)
        random.seed(42)
        np.random.seed(42)

        registration_matrix = np.array(
            [
                [1, 0, 0, 10],
                [0, 1, 0, 0],
                [0, 0, 1.2, 0],
                [0, 0, 0, 1],
            ]
        )

        subject_a = tio.Subject(
            t1=tio.ScalarImage(self.get_image_path('t1_a')),
        )
        subject_b = tio.Subject(
            t1=tio.ScalarImage(self.get_image_path('t1_b')),
            label=tio.LabelMap(self.get_image_path('label_b', binary=True)),
        )
        subject_c = tio.Subject(
            label=tio.LabelMap(self.get_image_path('label_c', binary=True)),
        )
        subject_d = tio.Subject(
            t1=tio.ScalarImage(
                self.get_image_path('t1_d'),
                pre_affine=registration_matrix,
            ),
            t2=tio.ScalarImage(self.get_image_path('t2_d')),
            label=tio.LabelMap(self.get_image_path('label_d', binary=True)),
        )
        subject_a4 = tio.Subject(
            t1=tio.ScalarImage(self.get_image_path('t1_a'), components=4),
        )
        self.subjects_list = [
            subject_a,
            subject_a4,
            subject_b,
            subject_c,
            subject_d,
        ]
        self.dataset = tio.SubjectsDataset(self.subjects_list)
        self.sample_subject = self.dataset[-1]  # subject_d
        self.subject_4d = self.dataset[1]

    def make_2d(self, subject):
        subject = copy.deepcopy(subject)
        for image in subject.get_images(intensity_only=False):
            image.set_data(image.data[..., :1])
        return subject

    def make_multichannel(self, subject):
        subject = copy.deepcopy(subject)
        for image in subject.get_images(intensity_only=False):
            image.set_data(torch.cat(4 * (image.data,)))
        return subject

    def flip_affine_x(self, subject):
        subject = copy.deepcopy(subject)
        for image in subject.get_images(intensity_only=False):
            image.affine = np.diag((-1, 1, 1, 1)) @ image.affine
        return subject

    def get_inconsistent_shape_subject(self):
        """Return a subject containing images of different shape."""
        subject = tio.Subject(
            t1=tio.ScalarImage(self.get_image_path('t1_inc')),
            t2=tio.ScalarImage(
                self.get_image_path('t2_inc', shape=(10, 20, 31)),
            ),
            label=tio.LabelMap(
                self.get_image_path(
                    'label_inc',
                    shape=(8, 17, 25),
                    binary=True,
                ),
            ),
            label2=tio.LabelMap(
                self.get_image_path(
                    'label2_inc',
                    shape=(18, 17, 25),
                    binary=True,
                ),
            ),
        )
        return subject

    def get_reference_image_and_path(self):
        """Return a reference image and its path."""
        path = self.get_image_path(
            'ref',
            shape=(10, 20, 31),
            spacing=(1, 1, 2),
        )
        image = tio.ScalarImage(path)
        return image, path

    def get_subject_with_partial_volume_label_map(self, components=1):
        """Return a subject with a partial-volume label map."""
        return tio.Subject(
            t1=tio.ScalarImage(
                self.get_image_path('t1_d'),
            ),
            label=tio.LabelMap(
                self.get_image_path(
                    'label_d2',
                    binary=False,
                    components=components,
                ),
            ),
        )

    def get_subject_with_labels(self, labels):
        return tio.Subject(
            label=tio.LabelMap(
                self.get_image_path(
                    'label_multi',
                    labels=labels,
                ),
            ),
        )

    @staticmethod
    def get_unique_labels(data: torch.Tensor) -> set[int]:
        labels = data.unique().tolist()
        return set(labels)

    @staticmethod
    def get_tensor_with_labels(labels: Sequence) -> torch.Tensor:
        tensor = torch.as_tensor(list(labels))
        return tensor.repeat_interleave(2).reshape(1, 1, 1, -1)

    def tearDown(self):
        """Tear down test fixtures, if any."""
        shutil.rmtree(self.dir)

    def get_ixi_tiny(self):
        root_dir = Path(tempfile.gettempdir()) / 'torchio' / 'ixi_tiny'
        return tio.datasets.IXITiny(root_dir, download=True)

    def get_image_path(
        self,
        stem,
        binary=False,
        labels=None,
        shape=(10, 20, 30),
        spacing=(1, 1, 1),
        components=1,
        add_nans=False,
        suffix=None,
        force_binary_foreground=True,
    ):
        shape = (*shape, 1) if len(shape) == 2 else shape
        data = np.random.rand(components, *shape)
        if binary:
            data = (data > 0.5).astype(np.uint8)
            if not data.sum() and force_binary_foreground:
                data[..., 0] = 1
        elif labels is not None:
            data = (data * (len(labels) + 1)).astype(np.uint8)
            new_data = np.zeros_like(data)
            for i, label in enumerate(labels):
                new_data[data == (i + 1)] = label
                if not (new_data == label).sum():
                    new_data[..., i] = label
            data = new_data
        elif self.flip_coin():  # cast some images
            data *= 100
            dtype = np.uint8 if self.flip_coin() else np.uint16
            data = data.astype(dtype)
        if add_nans:
            data[:] = np.nan
        affine = np.diag((*spacing, 1))
        if suffix is None:
            extensions = '.nii.gz', '.nii', '.nrrd', '.img', '.mnc'
            suffix = random.choice(extensions)
        path = self.dir / f'{stem}{suffix}'
        if self.flip_coin():
            path = str(path)
        image = tio.ScalarImage(
            tensor=data,
            affine=affine,
            check_nans=not add_nans,
        )
        image.save(path)
        return path

    def flip_coin(self):
        return np.random.rand() > 0.5

    def get_tests_data_dir(self):
        return Path(__file__).parent / 'image_data'

    def assert_tensor_not_equal(self, *args, **kwargs):  # noqa: N802
        with pytest.raises(AssertionError):
            self.assert_tensor_equal(*args, **kwargs)

    @staticmethod
    def assert_tensor_equal(*args, **kwargs):  # noqa: N802
        torch.testing.assert_close(
            *args,
            rtol=0,
            atol=0,
            check_dtype=False,
            **kwargs,
        )

    @staticmethod
    def assert_tensor_almost_equal(*args, **kwargs):  # noqa: N802
        torch.testing.assert_close(
            *args,
            **kwargs,
            check_dtype=False,
        )

    @staticmethod
    def assert_tensor_all_zeros(tensor):  # noqa: N802
        assert torch.all(tensor == 0)

    def get_large_composed_transform(self):
        all_classes = get_all_random_transforms()
        shuffle(all_classes)
        transforms = [t() for t in all_classes]
        # Hack as default patch size for RandomSwap is 15 and sample_subject
        # is (10, 20, 30)
        for tr in transforms:
            if tr.name == 'RandomSwap':
                tr.patch_size = np.array((10, 10, 10))
        return tio.Compose(transforms)


def get_all_random_transforms():
    transforms_names = [
        name for name in dir(tio.transforms) if name.startswith('Random')
    ]
    classes = [getattr(tio.transforms, name) for name in transforms_names]
    return classes