torchio.transforms.transform.DataToSubject._parse_subject() - Code Metrics - Inspection of "Add some deterministic transforms" - fepegar/torchio - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — master (#353)

by Fernando

created 2020-11-18 00:30 UTC

DataToSubject._parse_subject() A

↳ Parent: torchio.transforms.transform

Complexity

Conditions

Size

Total Lines	8
Code Lines	6

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	2
eloc	6
nop	1
dl	0
loc	8
rs	10
c	0
b	0
f	0

import copy
import numbers
from abc import ABC, abstractmethod
from typing import Optional, Union, Tuple, List

import torch
import numpy as np
import nibabel as nib
import SimpleITK as sitk

from .. import TypeData, DATA, AFFINE, TypeNumber
from ..data.subject import Subject
from ..data.image import Image, ScalarImage
from ..utils import nib_to_sitk, sitk_to_nib, to_tuple
from .interpolation import Interpolation


TypeTransformInput = Union[
    Subject,
    Image,
    torch.Tensor,
    np.ndarray,
    sitk.Image,
    dict,
    nib.Nifti1Image,
]


class Transform(ABC):
    """Abstract class for all TorchIO transforms.

    All subclasses should overwrite
    :py:meth:`torchio.tranforms.Transform.apply_transform`,
    which takes data, applies some transformation and returns the result.

    The input can be an instance of
    :py:class:`torchio.Subject`,
    :py:class:`torchio.Image`,
    :py:class:`numpy.ndarray`,
    :py:class:`torch.Tensor`,
    :py:class:`SimpleITK.image`,
    or a Python dictionary.

    Args:
        p: Probability that this transform will be applied.
        copy: Make a shallow copy of the input before applying the transform.
        keys: Mandatory if the input is a Python dictionary. The transform will
            be applied only to the data in each key.
    """
    def __init__(
            self,
            p: float = 1,
            copy: bool = True,
            keys: Optional[List[str]] = None,
            ):
        self.probability = self.parse_probability(p)
        self.copy = copy
        self.keys = keys

    def __call__(
            self,
            data: TypeTransformInput,
            ) -> TypeTransformInput:
        """Transform data and return a result of the same type.

        Args:
            data: Instance of 1) :py:class:`~torchio.Subject`, 4D
                :py:class:`torch.Tensor` or NumPy array with dimensions
                :math:`(C, W, H, D)`, where :math:`C` is the number of channels
                and :math:`W, H, D` are the spatial dimensions. If the input is
                a tensor, the affine matrix will be set to identity. Other
                valid input types are a SimpleITK image, a
                :py:class:`torch.Image`, a NiBabel Nifti1 Image or a Python
                dictionary. The output type is the same as te input type.
        """
        if torch.rand(1).item() > self.probability:
            return data
        data_parser = DataToSubject(data, keys=self.keys)
        subject = data_parser.get_subject()
        if self.copy:
            subject = copy.copy(subject)
        with np.errstate(all='raise'):
            transformed = self.apply_transform(subject)
        self.add_transform_to_subject_history(transformed)
        for image in transformed.get_images(intensity_only=False):
            ndim = image[DATA].ndim
            assert ndim == 4, f'Output of {self.name} is {ndim}D'
        output = data_parser.get_output(transformed)
        return output

    def __repr__(self):
        if hasattr(self, 'args_names'):
            names = self.args_names
            args_strings = (f'{arg}={getattr(self, arg)}' for arg in names)
            args_string = ', '.join(args_strings)
            return f'{self.name}({args_string})'
        else:
            return super().__repr__()

    @abstractmethod
    def apply_transform(self, subject: Subject):
        raise NotImplementedError

    def add_transform_to_subject_history(self, subject):
        from .augmentation import RandomTransform
        from . import Compose, OneOf, CropOrPad
        call_others = (
            RandomTransform,
            Compose,
            OneOf,
            CropOrPad,
        )
        if not isinstance(self, call_others):
            subject.add_transform(self, self.get_arguments())

    @staticmethod
    def to_range(n, around):
        if around is None:
            return 0, n
        else:
            return around - n, around + n

    def parse_params(self, params, around, name, make_ranges=True, **kwargs):
        params = to_tuple(params)
        if len(params) == 1 or (len(params) == 2 and make_ranges):  # d or (a, b)
            params *= 3  # (d, d, d) or (a, b, a, b, a, b)
        if len(params) == 3 and make_ranges:  # (a, b, c)
            items = [self.to_range(n, around) for n in params]
            # (-a, a, -b, b, -c, c) or (1-a, 1+a, 1-b, 1+b, 1-c, 1+c)
            params = [n for prange in items for n in prange]
        if make_ranges and len(params) != 6:
            if len(params) != 6:
                message = (
                    f'If "{name}" is a sequence, it must have length 2, 3 or 6,'
                    f' not {len(params)}'
                )
                raise ValueError(message)
            for param_range in zip(params[::2], params[1::2]):
                self.parse_range(param_range, name, **kwargs)
        return tuple(params)

    @staticmethod
    def parse_range(
            nums_range: Union[TypeNumber, Tuple[TypeNumber, TypeNumber]],
            name: str,
            min_constraint: TypeNumber = None,
            max_constraint: TypeNumber = None,
            type_constraint: type = None,
            ) -> Tuple[TypeNumber, TypeNumber]:
        r"""Adapted from ``torchvision.transforms.RandomRotation``.

        Args:
            nums_range: Tuple of two numbers :math:`(n_{min}, n_{max})`,
                where :math:`n_{min} \leq n_{max}`.
                If a single positive number :math:`n` is provided,
                :math:`n_{min} = -n` and :math:`n_{max} = n`.
            name: Name of the parameter, so that an informative error message
                can be printed.
            min_constraint: Minimal value that :math:`n_{min}` can take,
                default is None, i.e. there is no minimal value.
            max_constraint: Maximal value that :math:`n_{max}` can take,
                default is None, i.e. there is no maximal value.
            type_constraint: Precise type that :math:`n_{max}` and
                :math:`n_{min}` must take.

        Returns:
            A tuple of two numbers :math:`(n_{min}, n_{max})`.

        Raises:
            ValueError: if :attr:`nums_range` is negative
            ValueError: if :math:`n_{max}` or :math:`n_{min}` is not a number
            ValueError: if :math:`n_{max} \lt n_{min}`
            ValueError: if :attr:`min_constraint` is not None and
                :math:`n_{min}` is smaller than :attr:`min_constraint`
            ValueError: if :attr:`max_constraint` is not None and
                :math:`n_{max}` is greater than :attr:`max_constraint`
            ValueError: if :attr:`type_constraint` is not None and
                :math:`n_{max}` and :math:`n_{max}` are not of type
                :attr:`type_constraint`.
        """
        if isinstance(nums_range, numbers.Number):  # single number given
            if nums_range < 0:
                raise ValueError(
                    f'If {name} is a single number,'
                    f' it must be positive, not {nums_range}')
            if min_constraint is not None and nums_range < min_constraint:
                raise ValueError(
                    f'If {name} is a single number, it must be greater'
                    f' than {min_constraint}, not {nums_range}'
                )
            if max_constraint is not None and nums_range > max_constraint:
                raise ValueError(
                    f'If {name} is a single number, it must be smaller'
                    f' than {max_constraint}, not {nums_range}'
                )
            if type_constraint is not None:
                if not isinstance(nums_range, type_constraint):
                    raise ValueError(
                        f'If {name} is a single number, it must be of'
                        f' type {type_constraint}, not {nums_range}'
                    )
            min_range = -nums_range if min_constraint is None else nums_range
            return (min_range, nums_range)

        try:
            min_value, max_value = nums_range
        except (TypeError, ValueError):
            raise ValueError(
                f'If {name} is not a single number, it must be'
                f' a sequence of len 2, not {nums_range}'
            )

        min_is_number = isinstance(min_value, numbers.Number)
        max_is_number = isinstance(max_value, numbers.Number)
        if not min_is_number or not max_is_number:
            message = (
                f'{name} values must be numbers, not {nums_range}')
            raise ValueError(message)

        if min_value > max_value:
            raise ValueError(
                f'If {name} is a sequence, the second value must be'
                f' equal or greater than the first, but it is {nums_range}')

        if min_constraint is not None and min_value < min_constraint:
            raise ValueError(
                f'If {name} is a sequence, the first value must be greater'
                f' than {min_constraint}, but it is {min_value}'
            )

        if max_constraint is not None and max_value > max_constraint:
            raise ValueError(
                f'If {name} is a sequence, the second value must be smaller'
                f' than {max_constraint}, but it is {max_value}'
            )

        if type_constraint is not None:
            min_type_ok = isinstance(min_value, type_constraint)
            max_type_ok = isinstance(max_value, type_constraint)
            if not min_type_ok or not max_type_ok:
                raise ValueError(
                    f'If "{name}" is a sequence, its values must be of'
                    f' type "{type_constraint}", not "{type(nums_range)}"'
                )
        return nums_range

    @staticmethod
    def parse_interpolation(interpolation: str) -> str:
        interpolation = interpolation.lower()
        is_string = isinstance(interpolation, str)
        supported_values = [key.name.lower() for key in Interpolation]
        is_supported = interpolation.lower() in supported_values
        if is_string and is_supported:
            return interpolation
        message = (
            f'Interpolation "{interpolation}" of type {type(interpolation)}'
            f' must be a string among the supported values: {supported_values}'
        )
        raise TypeError(message)

    @staticmethod
    def parse_probability(probability: float) -> float:
        is_number = isinstance(probability, numbers.Number)
        if not (is_number and 0 <= probability <= 1):
            message = (
                'Probability must be a number in [0, 1],'
                f' not {probability}'
            )
            raise ValueError(message)
        return probability

    @staticmethod
    def nib_to_sitk(data: TypeData, affine: TypeData) -> sitk.Image:
        return nib_to_sitk(data, affine)

    @staticmethod
    def sitk_to_nib(image: sitk.Image) -> Tuple[torch.Tensor, np.ndarray]:
        return sitk_to_nib(image)

    @property
    def name(self):
        return self.__class__.__name__

    def get_arguments(self):
        """
        Return a dictionary with the arguments that would be necessary to
        reproduce the transform exactly.
        """
        return {name: getattr(self, name) for name in self.args_names}

    def is_invertible(self):
        return hasattr(self, 'invert_transform')

    def inverse(self):
        if not self.is_invertible():
            raise RuntimeError(f'{self.name} is not invertible')
        new = copy.deepcopy(self)
        new.invert_transform = not self.invert_transform
        return new


class DataToSubject:
    def __init__(
            self,
            data: TypeTransformInput,
            keys: Optional[List[str]] = None,
            ):
        self.data = data
        self.keys = keys
        self.default_image_name = 'default_image_name'
        self.is_tensor = False
        self.is_array = False
        self.is_dict = False
        self.is_image = False
        self.is_sitk = False
        self.is_nib = False

    def get_subject(self):
        if isinstance(self.data, nib.Nifti1Image):
            tensor = self.data.get_fdata(dtype=np.float32)
            data = ScalarImage(tensor=tensor, affine=self.data.affine)
            subject = self._get_subject_from_image(data)
            self.is_nib = True
        elif isinstance(self.data, (np.ndarray, torch.Tensor)):
            subject = self._parse_tensor(self.data)
            self.is_array = isinstance(self.data, np.ndarray)
            self.is_tensor = True
        elif isinstance(self.data, Image):
            subject = self._get_subject_from_image(self.data)
            self.is_image = True
        elif isinstance(self.data, Subject):
            subject = self.data
        elif isinstance(self.data, sitk.Image):
            subject = self._get_subject_from_sitk_image(self.data)
            self.is_sitk = True
        elif isinstance(self.data, dict):  # e.g. Eisen or MONAI dicts
            if self.keys is None:
                message = (
                    'If input is a dictionary, a value for "keys" must be'
                    ' specified when instantiating the transform'
                )
                raise RuntimeError(message)
            subject = self._get_subject_from_dict(self.data, self.keys)
            self.is_dict = True
        else:
            raise ValueError(f'Input type not recognized: {type(self.data)}')
        self._parse_subject(subject)
        return subject

    def get_output(self, transformed):
        if self.is_tensor or self.is_sitk:
            image = transformed[self.default_image_name]
            transformed = image[DATA]
            if self.is_array:
                transformed = transformed.numpy()
            elif self.is_sitk:
                transformed = nib_to_sitk(image[DATA], image[AFFINE])
        elif self.is_image:
            transformed = transformed[self.default_image_name]
        elif self.is_dict:
            transformed = dict(transformed)
            for key, value in transformed.items():
                if isinstance(value, Image):
                    transformed[key] = value.data
        elif self.is_nib:
            image = transformed[self.default_image_name]
            data = image[DATA]
            if len(data) > 1:
                message = (
                    'Multichannel images not supported for input of type'
                    ' nibabel.nifti.Nifti1Image'
                )
                raise RuntimeError(message)
            transformed = nib.Nifti1Image(data[0].numpy(), image[AFFINE])
        return transformed

    @staticmethod
    def _parse_subject(subject: Subject) -> None:
        if not isinstance(subject, Subject):
            message = (
                'Input to a transform must be a tensor or an instance'
                f' of torchio.Subject, not "{type(subject)}"'
            )
            raise RuntimeError(message)

    def _parse_tensor(self, data: TypeData) -> Subject:
        if data.ndim != 4:
            message = (
                'The input must be a 4D tensor with dimensions'
                f' (channels, x, y, z) but it has shape {tuple(data.shape)}'
            )
            raise ValueError(message)
        return self._get_subject_from_tensor(data)

    def _get_subject_from_tensor(self, tensor: torch.Tensor) -> Subject:
        image = ScalarImage(tensor=tensor)
        return self._get_subject_from_image(image)

    def _get_subject_from_image(self, image: Image) -> Subject:
        subject = Subject({self.default_image_name: image})
        return subject

    @staticmethod
    def _get_subject_from_dict(
            data: dict,
            image_keys: List[str],
            ) -> Subject:
        subject_dict = {}
        for key, value in data.items():
            if key in image_keys:
                value = ScalarImage(tensor=value)
            subject_dict[key] = value
        return Subject(subject_dict)

    def _get_subject_from_sitk_image(self, image):
        tensor, affine = sitk_to_nib(image)
        image = ScalarImage(tensor=tensor, affine=affine)
        return self._get_subject_from_image(image)


1			import copy
2			import numbers
3			from abc import ABC, abstractmethod
4			from typing import Optional, Union, Tuple, List
5
6			import torch
7			import numpy as np
8			import nibabel as nib
9			import SimpleITK as sitk
10
11			from .. import TypeData, DATA, AFFINE, TypeNumber
12			from ..data.subject import Subject
13			from ..data.image import Image, ScalarImage
14			from ..utils import nib_to_sitk, sitk_to_nib, to_tuple
15			from .interpolation import Interpolation
16
17
18			TypeTransformInput = Union[
19			Subject,
20			Image,
21			torch.Tensor,
22			np.ndarray,
23			sitk.Image,
24			dict,
25			nib.Nifti1Image,
26			]
27
28
29			class Transform(ABC):
30			"""Abstract class for all TorchIO transforms.
31
32			All subclasses should overwrite
33			:py:meth:`torchio.tranforms.Transform.apply_transform`,
34			which takes data, applies some transformation and returns the result.
35
36			The input can be an instance of
37			:py:class:`torchio.Subject`,
38			:py:class:`torchio.Image`,
39			:py:class:`numpy.ndarray`,
40			:py:class:`torch.Tensor`,
41			:py:class:`SimpleITK.image`,
42			or a Python dictionary.
43
44			Args:
45			p: Probability that this transform will be applied.
46			copy: Make a shallow copy of the input before applying the transform.
47			keys: Mandatory if the input is a Python dictionary. The transform will
48			be applied only to the data in each key.
49			"""
50			def __init__(
51			self,
52			p: float = 1,
53			copy: bool = True,
54			keys: Optional[List[str]] = None,
55			):
56			self.probability = self.parse_probability(p)
57			self.copy = copy
58			self.keys = keys
59
60			def __call__(
61			self,
62			data: TypeTransformInput,
63			) -> TypeTransformInput:
64			"""Transform data and return a result of the same type.
65
66			Args:
67			data: Instance of 1) :py:class:`~torchio.Subject`, 4D
68			:py:class:`torch.Tensor` or NumPy array with dimensions
69			:math:`(C, W, H, D)`, where :math:`C` is the number of channels
70			and :math:`W, H, D` are the spatial dimensions. If the input is
71			a tensor, the affine matrix will be set to identity. Other
72			valid input types are a SimpleITK image, a
73			:py:class:`torch.Image`, a NiBabel Nifti1 Image or a Python
74			dictionary. The output type is the same as te input type.
75			"""
76			if torch.rand(1).item() > self.probability:
77			return data
78			data_parser = DataToSubject(data, keys=self.keys)
79			subject = data_parser.get_subject()
80			if self.copy:
81			subject = copy.copy(subject)
82			with np.errstate(all='raise'):
83			transformed = self.apply_transform(subject)
84			self.add_transform_to_subject_history(transformed)
85			for image in transformed.get_images(intensity_only=False):
86			ndim = image[DATA].ndim
87			assert ndim == 4, f'Output of {self.name} is {ndim}D'
88			output = data_parser.get_output(transformed)
89			return output
90
91			def __repr__(self):
92			if hasattr(self, 'args_names'):
93			names = self.args_names
94			args_strings = (f'{arg}={getattr(self, arg)}' for arg in names)
95			args_string = ', '.join(args_strings)
96			return f'{self.name}({args_string})'
97			else:
98			return super().__repr__()
99
100			@abstractmethod
101			def apply_transform(self, subject: Subject):
102			raise NotImplementedError
103
104			def add_transform_to_subject_history(self, subject):
105			from .augmentation import RandomTransform
106			from . import Compose, OneOf, CropOrPad
107			call_others = (
108			RandomTransform,
109			Compose,
110			OneOf,
111			CropOrPad,
112			)
113			if not isinstance(self, call_others):
114			subject.add_transform(self, self.get_arguments())
115
116			@staticmethod
117			def to_range(n, around):
118			if around is None:
119			return 0, n
120			else:
121			return around - n, around + n
122
123			def parse_params(self, params, around, name, make_ranges=True, **kwargs):
124			params = to_tuple(params)
125			if len(params) == 1 or (len(params) == 2 and make_ranges): # d or (a, b)
126			params *= 3 # (d, d, d) or (a, b, a, b, a, b)
127			if len(params) == 3 and make_ranges: # (a, b, c)
128			items = [self.to_range(n, around) for n in params]
129			# (-a, a, -b, b, -c, c) or (1-a, 1+a, 1-b, 1+b, 1-c, 1+c)
130			params = [n for prange in items for n in prange]
131			if make_ranges and len(params) != 6:
132			if len(params) != 6:
133			message = (
134			f'If "{name}" is a sequence, it must have length 2, 3 or 6,'
135			f' not {len(params)}'
136			)
137			raise ValueError(message)
138			for param_range in zip(params[::2], params[1::2]):
139			self.parse_range(param_range, name, **kwargs)
140			return tuple(params)
141
142			@staticmethod
143			def parse_range(
144			nums_range: Union[TypeNumber, Tuple[TypeNumber, TypeNumber]],
145			name: str,
146			min_constraint: TypeNumber = None,
147			max_constraint: TypeNumber = None,
148			type_constraint: type = None,
149			) -> Tuple[TypeNumber, TypeNumber]:
150			r"""Adapted from ``torchvision.transforms.RandomRotation``.
151
152			Args:
153			nums_range: Tuple of two numbers :math:`(n_{min}, n_{max})`,
154			where :math:`n_{min} \leq n_{max}`.
155			If a single positive number :math:`n` is provided,
156			:math:`n_{min} = -n` and :math:`n_{max} = n`.
157			name: Name of the parameter, so that an informative error message
158			can be printed.
159			min_constraint: Minimal value that :math:`n_{min}` can take,
160			default is None, i.e. there is no minimal value.
161			max_constraint: Maximal value that :math:`n_{max}` can take,
162			default is None, i.e. there is no maximal value.
163			type_constraint: Precise type that :math:`n_{max}` and
164			:math:`n_{min}` must take.
165
166			Returns:
167			A tuple of two numbers :math:`(n_{min}, n_{max})`.
168
169			Raises:
170			ValueError: if :attr:`nums_range` is negative
171			ValueError: if :math:`n_{max}` or :math:`n_{min}` is not a number
172			ValueError: if :math:`n_{max} \lt n_{min}`
173			ValueError: if :attr:`min_constraint` is not None and
174			:math:`n_{min}` is smaller than :attr:`min_constraint`
175			ValueError: if :attr:`max_constraint` is not None and
176			:math:`n_{max}` is greater than :attr:`max_constraint`
177			ValueError: if :attr:`type_constraint` is not None and
178			:math:`n_{max}` and :math:`n_{max}` are not of type
179			:attr:`type_constraint`.
180			"""
181			if isinstance(nums_range, numbers.Number): # single number given
182			if nums_range < 0:
183			raise ValueError(
184			f'If {name} is a single number,'
185			f' it must be positive, not {nums_range}')
186			if min_constraint is not None and nums_range < min_constraint:
187			raise ValueError(
188			f'If {name} is a single number, it must be greater'
189			f' than {min_constraint}, not {nums_range}'
190			)
191			if max_constraint is not None and nums_range > max_constraint:
192			raise ValueError(
193			f'If {name} is a single number, it must be smaller'
194			f' than {max_constraint}, not {nums_range}'
195			)
196			if type_constraint is not None:
197			if not isinstance(nums_range, type_constraint):
198			raise ValueError(
199			f'If {name} is a single number, it must be of'
200			f' type {type_constraint}, not {nums_range}'
201			)
202			min_range = -nums_range if min_constraint is None else nums_range
203			return (min_range, nums_range)
204
205			try:
206			min_value, max_value = nums_range
207			except (TypeError, ValueError):
208			raise ValueError(
209			f'If {name} is not a single number, it must be'
210			f' a sequence of len 2, not {nums_range}'
211			)
212
213			min_is_number = isinstance(min_value, numbers.Number)
214			max_is_number = isinstance(max_value, numbers.Number)
215			if not min_is_number or not max_is_number:
216			message = (
217			f'{name} values must be numbers, not {nums_range}')
218			raise ValueError(message)
219
220			if min_value > max_value:
221			raise ValueError(
222			f'If {name} is a sequence, the second value must be'
223			f' equal or greater than the first, but it is {nums_range}')
224
225			if min_constraint is not None and min_value < min_constraint:
226			raise ValueError(
227			f'If {name} is a sequence, the first value must be greater'
228			f' than {min_constraint}, but it is {min_value}'
229			)
230
231			if max_constraint is not None and max_value > max_constraint:
232			raise ValueError(
233			f'If {name} is a sequence, the second value must be smaller'
234			f' than {max_constraint}, but it is {max_value}'
235			)
236
237			if type_constraint is not None:
238			min_type_ok = isinstance(min_value, type_constraint)
239			max_type_ok = isinstance(max_value, type_constraint)
240			if not min_type_ok or not max_type_ok:
241			raise ValueError(
242			f'If "{name}" is a sequence, its values must be of'
243			f' type "{type_constraint}", not "{type(nums_range)}"'
244			)
245			return nums_range
246
247			@staticmethod
248			def parse_interpolation(interpolation: str) -> str:
249			interpolation = interpolation.lower()
250			is_string = isinstance(interpolation, str)
251			supported_values = [key.name.lower() for key in Interpolation]
252			is_supported = interpolation.lower() in supported_values
253			if is_string and is_supported:
254			return interpolation
255			message = (
256			f'Interpolation "{interpolation}" of type {type(interpolation)}'
257			f' must be a string among the supported values: {supported_values}'
258			)
259			raise TypeError(message)
260
261			@staticmethod
262			def parse_probability(probability: float) -> float:
263			is_number = isinstance(probability, numbers.Number)
264			if not (is_number and 0 <= probability <= 1):
265			message = (
266			'Probability must be a number in [0, 1],'
267			f' not {probability}'
268			)
269			raise ValueError(message)
270			return probability
271
272			@staticmethod
273			def nib_to_sitk(data: TypeData, affine: TypeData) -> sitk.Image:
274			return nib_to_sitk(data, affine)
275
276			@staticmethod
277			def sitk_to_nib(image: sitk.Image) -> Tuple[torch.Tensor, np.ndarray]:
278			return sitk_to_nib(image)
279
280			@property
281			def name(self):
282			return self.__class__.__name__
283
284			def get_arguments(self):
285			"""
286			Return a dictionary with the arguments that would be necessary to
287			reproduce the transform exactly.
288			"""
289			return {name: getattr(self, name) for name in self.args_names}
290
291			def is_invertible(self):
292			return hasattr(self, 'invert_transform')
293
294			def inverse(self):
295			if not self.is_invertible():
296			raise RuntimeError(f'{self.name} is not invertible')
297			new = copy.deepcopy(self)
298			new.invert_transform = not self.invert_transform
299			return new
300
301
302			class DataToSubject:
303			def __init__(
304			self,
305			data: TypeTransformInput,
306			keys: Optional[List[str]] = None,
307			):
308			self.data = data
309			self.keys = keys
310			self.default_image_name = 'default_image_name'
311			self.is_tensor = False
312			self.is_array = False
313			self.is_dict = False
314			self.is_image = False
315			self.is_sitk = False
316			self.is_nib = False
317
318			def get_subject(self):
319			if isinstance(self.data, nib.Nifti1Image):
320			tensor = self.data.get_fdata(dtype=np.float32)
321			data = ScalarImage(tensor=tensor, affine=self.data.affine)
322			subject = self._get_subject_from_image(data)
323			self.is_nib = True
324			elif isinstance(self.data, (np.ndarray, torch.Tensor)):
325			subject = self._parse_tensor(self.data)
326			self.is_array = isinstance(self.data, np.ndarray)
327			self.is_tensor = True
328			elif isinstance(self.data, Image):
329			subject = self._get_subject_from_image(self.data)
330			self.is_image = True
331			elif isinstance(self.data, Subject):
332			subject = self.data
333			elif isinstance(self.data, sitk.Image):
334			subject = self._get_subject_from_sitk_image(self.data)
335			self.is_sitk = True
336			elif isinstance(self.data, dict): # e.g. Eisen or MONAI dicts
337			if self.keys is None:
338			message = (
339			'If input is a dictionary, a value for "keys" must be'
340			' specified when instantiating the transform'
341			)
342			raise RuntimeError(message)
343			subject = self._get_subject_from_dict(self.data, self.keys)
344			self.is_dict = True
345			else:
346			raise ValueError(f'Input type not recognized: {type(self.data)}')
347			self._parse_subject(subject)
348			return subject
349
350			def get_output(self, transformed):
351			if self.is_tensor or self.is_sitk:
352			image = transformed[self.default_image_name]
353			transformed = image[DATA]
354			if self.is_array:
355			transformed = transformed.numpy()
356			elif self.is_sitk:
357			transformed = nib_to_sitk(image[DATA], image[AFFINE])
358			elif self.is_image:
359			transformed = transformed[self.default_image_name]
360			elif self.is_dict:
361			transformed = dict(transformed)
362			for key, value in transformed.items():
363			if isinstance(value, Image):
364			transformed[key] = value.data
365			elif self.is_nib:
366			image = transformed[self.default_image_name]
367			data = image[DATA]
368			if len(data) > 1:
369			message = (
370			'Multichannel images not supported for input of type'
371			' nibabel.nifti.Nifti1Image'
372			)
373			raise RuntimeError(message)
374			transformed = nib.Nifti1Image(data[0].numpy(), image[AFFINE])
375			return transformed
376
377			@staticmethod
378			def _parse_subject(subject: Subject) -> None:
379			if not isinstance(subject, Subject):
380			message = (
381			'Input to a transform must be a tensor or an instance'
382			f' of torchio.Subject, not "{type(subject)}"'
383			)
384			raise RuntimeError(message)
385
386			def _parse_tensor(self, data: TypeData) -> Subject:
387			if data.ndim != 4:
388			message = (
389			'The input must be a 4D tensor with dimensions'
390			f' (channels, x, y, z) but it has shape {tuple(data.shape)}'
391			)
392			raise ValueError(message)
393			return self._get_subject_from_tensor(data)
394
395			def _get_subject_from_tensor(self, tensor: torch.Tensor) -> Subject:
396			image = ScalarImage(tensor=tensor)
397			return self._get_subject_from_image(image)
398
399			def _get_subject_from_image(self, image: Image) -> Subject:
400			subject = Subject({self.default_image_name: image})
401			return subject
402
403			@staticmethod
404			def _get_subject_from_dict(
405			data: dict,
406			image_keys: List[str],
407			) -> Subject:
408			subject_dict = {}
409			for key, value in data.items():
410			if key in image_keys:
411			value = ScalarImage(tensor=value)
412			subject_dict[key] = value
413			return Subject(subject_dict)
414
415			def _get_subject_from_sitk_image(self, image):
416			tensor, affine = sitk_to_nib(image)
417			image = ScalarImage(tensor=tensor, affine=affine)
418			return self._get_subject_from_image(image)
419

fepegar / torchio

Pull Request — master (#353)

DataToSubject._parse_subject() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like