Add Contour transform

fepegar · Jan 23, 2021 · ce3c5af · ce3c5af
1 parent 6e53c83
commit ce3c5af
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 4 deletions.
diff --git a/docs/source/transforms/preprocessing.rst b/docs/source/transforms/preprocessing.rst
@@ -122,6 +122,13 @@ Label
     :show-inheritance:
 
 
+:class:`Contour`
+~~~~~~~~~~~~~~~~
+
+.. autoclass:: Contour
+    :show-inheritance:
+
+
 :class:`KeepLargestComponent`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

diff --git a/torchio/transforms/__init__.py b/torchio/transforms/__init__.py
@@ -37,6 +37,7 @@
 from .preprocessing import HistogramStandardization
 from .preprocessing.intensity.histogram_standardization import train_histogram
 from .preprocessing import OneHot
+from .preprocessing import Contour
 from .preprocessing import RemapLabels
 from .preprocessing import RemoveLabels
 from .preprocessing import SequentialLabels
@@ -87,6 +88,7 @@
     'EnsureShapeMultiple',
     'train_histogram',
     'OneHot',
+    'Contour',
     'RemapLabels',
     'RemoveLabels',
     'SequentialLabels',

diff --git a/torchio/transforms/preprocessing/__init__.py b/torchio/transforms/preprocessing/__init__.py
@@ -10,6 +10,7 @@
 from .intensity.histogram_standardization import HistogramStandardization
 
 from .label.one_hot import OneHot
+from .label.contour import Contour
 from .label.remap_labels import RemapLabels
 from .label.remove_labels import RemoveLabels
 from .label.sequential_labels import SequentialLabels
@@ -27,6 +28,7 @@
     'RescaleIntensity',
     'HistogramStandardization',
     'OneHot',
+    'Contour',
     'RemapLabels',
     'RemoveLabels',
     'SequentialLabels',

diff --git a/torchio/transforms/preprocessing/label/contour.py b/torchio/transforms/preprocessing/label/contour.py
@@ -0,0 +1,24 @@
+import SimpleITK as sitk
+
+from .label_transform import LabelTransform
+
+
+class Contour(LabelTransform):
+    r"""Keep only the borders of each connected component in a binary image.
+
+    Args:
+        **kwargs: See :class:`~torchio.transforms.Transform` for additional
+            keyword arguments.
+    """
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.args_names = []
+
+    def apply_transform(self, subject):
+        for image in self.get_images(subject):
+            assert image.data.ndim == 4 and image.data.shape[0] == 1
+            sitk_image = image.as_sitk()
+            contour = sitk.BinaryContour(sitk_image)
+            tensor, _ = self.sitk_to_nib(contour)
+            image.set_data(tensor)
+        return subject
diff --git a/torchio/transforms/preprocessing/label/keep_largest_component.py b/torchio/transforms/preprocessing/label/keep_largest_component.py
@@ -16,10 +16,7 @@ class KeepLargestComponent(LabelTransform):
 
     .. _open a new issue: https://github.com/fepegar/torchio/issues/new?assignees=&labels=enhancement&template=feature_request.md&title=Improve%20KeepLargestComponent%20transform
     """  # noqa: E501
-    def __init__(
-            self,
-            **kwargs
-            ):
+    def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.args_names = []