Total Complexity | 3 |
Total Lines | 33 |
Duplicated Lines | 0 % |
Changes | 0 |
1 | import SimpleITK as sitk |
||
2 | |||
3 | from .label_transform import LabelTransform |
||
4 | |||
5 | |||
6 | class KeepLargestComponent(LabelTransform): |
||
7 | r"""Keep only the largest connected component in a binary label map. |
||
8 | |||
9 | Args: |
||
10 | **kwargs: See :class:`~torchio.transforms.Transform` for additional |
||
11 | keyword arguments. |
||
12 | |||
13 | .. note:: For now, this transform only works for binary images, i.e., label |
||
14 | maps with a background and a foreground class. If you are interested in |
||
15 | extending this transform `open a new issue`_. |
||
16 | |||
17 | .. _open a new issue: https://github.com/fepegar/torchio/issues/new?assignees=&labels=enhancement&template=feature_request.md&title=Improve%20KeepLargestComponent%20transform |
||
18 | """ # noqa: E501 |
||
19 | def __init__(self, **kwargs): |
||
20 | super().__init__(**kwargs) |
||
21 | self.args_names = [] |
||
22 | |||
23 | def apply_transform(self, subject): |
||
24 | for image in self.get_images(subject): |
||
25 | assert image.data.ndim == 4 and image.data.shape[0] == 1 |
||
26 | sitk_image = image.as_sitk() |
||
27 | connected_components = sitk.ConnectedComponent(sitk_image) |
||
28 | labeled_cc = sitk.RelabelComponent(connected_components) |
||
29 | largest_cc = labeled_cc == 1 |
||
30 | tensor, _ = self.sitk_to_nib(largest_cc) |
||
31 | image.set_data(tensor) |
||
32 | return subject |
||
33 |