Passed
Pull Request — master (#139)
by Fernando
01:18
created

torchio.data.image.Image.parse_tensor()   A

Complexity

Conditions 3

Size

Total Lines 12
Code Lines 10

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 3
eloc 10
nop 1
dl 0
loc 12
rs 9.9
c 0
b 0
f 0
1
import warnings
2
from pathlib import Path
3
from typing import Any, Dict, Tuple, Optional
4
5
import torch
0 ignored issues
show
introduced by
Unable to import 'torch'
Loading history...
6
import numpy as np
0 ignored issues
show
introduced by
Unable to import 'numpy'
Loading history...
7
8
from ..torchio import TypePath, DATA, TYPE, AFFINE, PATH, STEM, INTENSITY
9
from .io import read_image
10
11
12
class Image(dict):
13
    r"""Class to store information about an image.
14
15
    Args:
16
        path: Path to a file that can be read by
17
            :mod:`SimpleITK` or :mod:`nibabel` or to a directory containing
18
            DICOM files.
19
        type: Type of image, such as :attr:`torchio.INTENSITY` or
20
            :attr:`torchio.LABEL`. This will be used by the transforms to
21
            decide whether to apply an operation, or which interpolation to use
22
            when resampling.
23
        tensor: If :attr:`path` is not given, :attr:`tensor` must be a 4D
24
            :py:class:`torch.Tensor` with dimensions :math:`(C, D, H, W)`,
25
            where :math:`C` is the number of channels and :math:`D, H, W`
26
            are the spatial dimensions.
27
        affine: If :attr:`path` is not given, :attr:`affine` must be a
28
            :math:`4 \times 4` NumPy array. If ``None``, :attr:`affine` is an
29
            identity matrix.
30
        **kwargs: Items that will be added to image dictionary within the
31
            subject sample.
32
    """
33
34
    def __init__(
35
            self,
36
            path: Optional[TypePath] = None,
37
            type: str = INTENSITY,
0 ignored issues
show
Bug Best Practice introduced by
This seems to re-define the built-in type.

It is generally discouraged to redefine built-ins as this makes code very hard to read.

Loading history...
38
            tensor: Optional[torch.Tensor] = None,
39
            affine: Optional[torch.Tensor] = None,
40
            **kwargs: Dict[str, Any],
41
            ):
42
        if path is None and tensor is None:
43
            raise ValueError('A value for path or tensor must be given')
44
        if path is not None:
45
            if tensor is not None or affine is not None:
46
                message = 'If a path is given, tensor and affine must be None'
47
                raise ValueError(message)
48
        self.tensor = self.parse_tensor(tensor)
49
        self.affine = self.parse_affine(affine)
50
        if self.affine is None:
51
            self.affine = np.eye(4)
52
        for key in (DATA, AFFINE, TYPE, PATH, STEM):
53
            if key in kwargs:
54
                raise ValueError(f'Key {key} is reserved. Use a different one')
55
56
        super().__init__(**kwargs)
57
        self.path = self._parse_path(path)
58
        self.type = type
59
        self.is_sample = False  # set to True by ImagesDataset
60
61
    @staticmethod
62
    def _parse_path(path: TypePath) -> Path:
0 ignored issues
show
Unused Code introduced by
Either all return statements in a function should return an expression, or none of them should.
Loading history...
63
        if path is None:
64
            return
65
        try:
66
            path = Path(path).expanduser()
67
        except TypeError:
68
            message = f'Conversion to path not possible for variable: {path}'
69
            raise TypeError(message)
70
        if not (path.is_file() or path.is_dir()):  # might be a dir with DICOM
71
            raise FileNotFoundError(f'File not found: {path}')
72
        return path
73
74
    @staticmethod
75
    def parse_tensor(tensor: torch.Tensor) -> torch.Tensor:
76
        if tensor is None:
77
            return None
78
        num_dimensions = tensor.dim()
79
        if num_dimensions != 4:
80
            message = (
81
                'The input tensor must have 4 dimensions (channels, i, j, k),'
82
                f' but has only {num_dimensions}: {tensor.shape}'
83
            )
84
            raise RuntimeError(message)
85
        return tensor
86
87
    @staticmethod
88
    def parse_affine(affine: np.ndarray) -> np.ndarray:
89
        if affine is None:
90
            return np.eye(4)
91
        if not isinstance(affine, np.ndarray):
92
            raise TypeError(f'Affine must be a NumPy array, not {type(affine)}')
93
        if affine.shape != (4, 4):
94
            raise ValueError(f'Affine shape must be (4, 4), not {affine.shape}')
95
        return affine
96
97
    def load(self, check_nans: bool = True) -> Tuple[torch.Tensor, np.ndarray]:
98
        r"""Load the image from disk.
99
100
        The file is expected to be monomodal and 3D. A channels dimension is
101
        added to the tensor.
102
103
        Args:
104
            check_nans: If ``True``, issues a warning if NaNs are found
105
                in the image
106
107
        Returns:
108
            Tuple containing a 4D data tensor of size
109
            :math:`(1, D_{in}, H_{in}, W_{in})`
110
            and a 2D 4x4 affine matrix
111
        """
112
        if self.path is None:
113
            return self.tensor, self.affine
114
        tensor, affine = read_image(self.path)
115
        tensor = tensor.unsqueeze(0)  # add channels dimension
116
        if check_nans and torch.isnan(tensor).any():
117
            warnings.warn(f'NaNs found in file "{self.path}"')
118
        return tensor, affine
119