Passed
Pull Request — main (#1308)
by
unknown
01:44
created

Crop._crop_image()   A

Complexity

Conditions 2

Size

Total Lines 22
Code Lines 17

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 2
eloc 17
nop 4
dl 0
loc 22
rs 9.55
c 0
b 0
f 0
1
from copy import deepcopy
2
3
import nibabel as nib
4
import numpy as np
5
6
from ....data.image import Image
7
from ....data.subject import Subject
8
from .bounds_transform import BoundsTransform
9
from .bounds_transform import TypeBounds
10
11
12
class Crop(BoundsTransform):
13
    r"""Crop an image.
14
15
    Args:
16
        cropping: Tuple
17
            :math:`(w_{ini}, w_{fin}, h_{ini}, h_{fin}, d_{ini}, d_{fin})`
18
            defining the number of values cropped from the edges of each axis.
19
            If the initial shape of the image is
20
            :math:`W \times H \times D`, the final shape will be
21
            :math:`(- w_{ini} + W - w_{fin}) \times (- h_{ini} + H - h_{fin})
22
            \times (- d_{ini} + D - d_{fin})`.
23
            If only three values :math:`(w, h, d)` are provided, then
24
            :math:`w_{ini} = w_{fin} = w`,
25
            :math:`h_{ini} = h_{fin} = h` and
26
            :math:`d_{ini} = d_{fin} = d`.
27
            If only one value :math:`n` is provided, then
28
            :math:`w_{ini} = w_{fin} = h_{ini} = h_{fin}
29
            = d_{ini} = d_{fin} = n`.
30
        copy: bool, optional
31
            This transform overwrites the copy argument of the base transform and
32
            copies only the cropped patch, instead of the whole image.
33
            This can provide a significant speedup when cropping small patches from large images
34
            If ``True``, each image will be cropped and the patch copied to a new subject.
35
            If ``False``, each image will be cropped in place. Default: ``True``.
36
        **kwargs: See :class:`~torchio.transforms.Transform` for additional
37
            keyword arguments.
38
39
    .. seealso:: If you want to pass the output shape instead, please use
40
        :class:`~torchio.transforms.CropOrPad` instead.
41
    """
42
43
    def __init__(self, cropping: TypeBounds, copy=True, **kwargs):
44
        self.copy_patch = copy
45
        # Transform base class deepcopies whole subject by default
46
        # We want to copy only the cropped patch, so we overwrite the functionality
47
        super().__init__(cropping, copy=False, **kwargs)
48
        self.cropping = cropping
49
        self.args_names = ['cropping']
50
51
    def apply_transform(self, subject: Subject) -> Subject:
52
        assert self.bounds_parameters is not None
53
        low = self.bounds_parameters[::2]
54
        high = self.bounds_parameters[1::2]
55
        index_ini = low
56
        index_fin = np.array(subject.spatial_shape) - high
57
58
        if self.copy_patch:
59
            # Create a new subject with only the cropped patch
60
            sample_attributes = {}
61
            image_keys_to_crop = subject.get_images_dict(
62
                intensity_only=False, include=self.include, exclude=self.exclude
63
            ).keys()
64
            # Copy all non-image attributes
65
            for key, value in subject.items():
66
                if key not in image_keys_to_crop:
67
                    sample_attributes[key] = deepcopy(value)
68
                else:
69
                    sample_attributes[key] = self._crop_image(
70
                        value, index_ini, index_fin
71
                    )
72
            cropped_sample = type(subject)(**sample_attributes)
73
74
            # Copy applied transforms history
75
            cropped_sample.applied_transforms = deepcopy(subject.applied_transforms)
76
77
            cropped_sample.update_attributes()
78
            return cropped_sample
79
        else:
80
            # Crop in place
81
            for image in self.get_images(subject):
82
                self._crop_image(image, index_ini, index_fin)
83
            return subject
84
85
    def _crop_image(self, image: Image, index_ini: tuple, index_fin: tuple) -> Image:
86
        new_origin = nib.affines.apply_affine(image.affine, index_ini)
87
        new_affine = image.affine.copy()
88
        new_affine[:3, 3] = new_origin
89
        i0, j0, k0 = index_ini
90
        i1, j1, k1 = index_fin
91
92
        # Crop the image data
93
        if self.copy_patch:
94
            # Create a new image with the cropped data
95
            cropped_data = image.data[:, i0:i1, j0:j1, k0:k1].clone()
96
            new_image = type(image)(
97
                tensor=cropped_data,
98
                affine=new_affine,
99
                type=image.type,
100
                path=image.path,
101
            )
102
            return new_image
103
        else:
104
            image.set_data(image.data[:, i0:i1, j0:j1, k0:k1].clone())
105
            image.affine = new_affine
106
            return image
107
108
    def inverse(self):
109
        from .pad import Pad
110
111
        return Pad(self.cropping)
112