Passed
Push — master ( 2e9b5a...e81273 )
by Fernando
01:23
created

Clamp.apply_clamp()   A

Complexity

Conditions 1

Size

Total Lines 2
Code Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 2
dl 0
loc 2
rs 10
c 0
b 0
f 0
cc 1
nop 2
1
from torchio.data.image import ScalarImage
2
from typing import Optional
3
import torch
4
from ....data.subject import Subject
5
from ...intensity_transform import IntensityTransform
6
7
8
class Clamp(IntensityTransform):
9
    """Clamp intensity values into a range :math:`[a, b]`.
10
11
    For more information, see :func:`torch.clamp`.
12
13
    Args:
14
        out_min: Minimum value :math:`a` of the output image. If ``None``, the
15
            minimum of the image is used.
16
        out_max: Maximum value :math:`b` of the output image. If ``None``, the
17
            maximum of the image is used.
18
19
    Example:
20
        >>> import torchio as tio
21
        >>> ct = tio.ScalarImage('ct_scan.nii.gz')
22
        >>> HOUNSFIELD_AIR, HOUNSFIELD_BONE = -1000, 1000
23
        >>> clamp = tio.Clamp(out_min=HOUNSFIELD_AIR, out_max=HOUNSFIELD_BONE)
24
        >>> ct_clamped = clamp(ct)
25
    """
26
    def __init__(
27
            self,
28
            out_min: Optional[float] = None,
29
            out_max: Optional[float] = None,
30
            **kwargs
31
            ):
32
        super().__init__(**kwargs)
33
        self.out_min, self.out_max = out_min, out_max
34
        self.args_names = 'out_min', 'out_max'
35
36
    def apply_transform(self, subject: Subject) -> Subject:
37
        for image in self.get_images(subject):
38
            self.apply_clamp(image)
39
        return subject
40
41
    def apply_clamp(self, image: ScalarImage) -> None:
42
        image.set_data(self.clamp(image.data))
43
44
    def clamp(self, tensor: torch.Tensor) -> torch.Tensor:
45
        return tensor.clamp(self.out_min, self.out_max)
46