Passed
Pull Request — master (#214)
by Fernando
01:47
created

torchio.data.subject.Subject.crop()   A

Complexity

Conditions 3

Size

Total Lines 10
Code Lines 8

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 3
eloc 8
nop 3
dl 0
loc 10
rs 10
c 0
b 0
f 0
1
import copy
2
import pprint
3
from typing import Any, Dict, List, Tuple
4
from ..torchio import TYPE, INTENSITY
5
from .image import Image
6
7
8
class Subject(dict):
9
    """Class to store information about the images corresponding to a subject.
10
11
    Args:
12
        *args: If provided, a dictionary of items.
13
        **kwargs: Items that will be added to the subject sample.
14
15
    Example:
16
17
        >>> import torchio
18
        >>> from torchio import Image, Subject
19
        >>> # One way:
20
        >>> subject = Subject(
21
        ...     one_image=Image('path_to_image.nii.gz', type=torchio.INTENSITY),
22
        ...     a_segmentation=Image('path_to_seg.nii.gz', type=torchio.LABEL),
23
        ...     age=45,
24
        ...     name='John Doe',
25
        ...     hospital='Hospital Juan Negrín',
26
        ... )
27
        >>> # If you want to create the mapping before, or have spaces in the keys:
28
        >>> subject_dict = {
29
        ...     'one image': Image('path_to_image.nii.gz', type=torchio.INTENSITY),
30
        ...     'a segmentation': Image('path_to_seg.nii.gz', type=torchio.LABEL),
31
        ...     'age': 45,
32
        ...     'name': 'John Doe',
33
        ...     'hospital': 'Hospital Juan Negrín',
34
        ... }
35
        >>> Subject(subject_dict)
36
37
    """
38
39
    def __init__(self, *args, **kwargs: Dict[str, Any]):
40
        if args:
41
            if len(args) == 1 and isinstance(args[0], dict):
42
                kwargs.update(args[0])
43
            else:
44
                message = (
45
                    'Only one dictionary as positional argument is allowed')
46
                raise ValueError(message)
47
        super().__init__(**kwargs)
48
        self.images = [
49
            (k, v) for (k, v) in self.items()
50
            if isinstance(v, Image)
51
        ]
52
        self._parse_images(self.images)
53
        self.__dict__.update(self)  # this allows me to do e.g. subject.t1
54
        self.history = []
55
56
    def __repr__(self):
57
        string = (
58
            f'{self.__class__.__name__}'
59
            f'(Keys: {tuple(self.keys())}; images: {len(self.images)})'
60
        )
61
        return string
62
63
    @staticmethod
64
    def _parse_images(images: List[Tuple[str, Image]]) -> None:
65
        # Check that it's not empty
66
        if not images:
67
            raise ValueError('A subject without images cannot be created')
68
69
    @property
70
    def shape(self):
71
        """Return shape of first image in sample.
72
73
        Consistency of shapes across images in the sample is checked first.
74
        """
75
        self.check_consistent_shape()
76
        image = self.get_images(intensity_only=False)[0]
77
        return image.shape
78
79
    @property
80
    def spatial_shape(self):
81
        """Return spatial shape of first image in sample.
82
83
        Consistency of shapes across images in the sample is checked first.
84
        """
85
        return self.shape[1:]
86
87
    def get_images_dict(self, intensity_only=True):
88
        images = {}
89
        for image_name, image in self.items():
90
            if not isinstance(image, Image):
91
                continue
92
            if intensity_only and not image[TYPE] == INTENSITY:
93
                continue
94
            images[image_name] = image
95
        return images
96
97
    def get_images(self, intensity_only=True):
98
        images_dict = self.get_images_dict(intensity_only=intensity_only)
99
        return list(images_dict.values())
100
101
    def check_consistent_shape(self) -> None:
102
        shapes_dict = {}
103
        iterable = self.get_images_dict(intensity_only=False).items()
104
        for image_name, image in iterable:
105
            shapes_dict[image_name] = image.shape
106
        num_unique_shapes = len(set(shapes_dict.values()))
107
        if num_unique_shapes > 1:
108
            message = (
109
                'Images in sample have inconsistent shapes:'
110
                f'\n{pprint.pformat(shapes_dict)}'
111
            )
112
            raise ValueError(message)
113
114
    def add_transform(
115
            self,
116
            transform: 'Transform',
117
            parameters_dict: dict,
118
            ) -> None:
119
        self.history.append((transform.name, parameters_dict))
120
121
    def load(self):
122
        for image in self.get_images(intensity_only=False):
123
            image.load()
124
125
    def crop(self, index_ini, index_fin):
126
        result_dict = {}
127
        for key, value in self.items():
128
            if isinstance(value, Image):
129
                # patch.clone() is much faster than copy.deepcopy(patch)
130
                value = value.crop(index_ini, index_fin)
131
            else:
132
                value = copy.deepcopy(value)
133
            result_dict[key] = value
134
        return Subject(result_dict)
135