SQUASHME: fix PR

Jo-Byr · Jo-Byr · commit de449480c500 · 2025-11-28T15:33:39.000+01:00
diff --git a/examples/viewer_lib/logic/segmentation/islands_effect_logic.py b/examples/viewer_lib/logic/segmentation/islands_effect_logic.py
@@ -26,21 +26,8 @@ def _on_apply_clicked(self):
         if not self.is_active():
             return
         if self._typed_state.data.mode == IslandsSegmentationMode.KEEP_LARGEST_ISLAND.value:
-            self.effect.apply(
-                segment_id=self.segmentation_editor.get_active_segment_id(),
-                max_number_of_segments=1,
-                split=False,
-            )
+            self.effect.keep_largest_island()
         elif self._typed_state.data.mode == IslandsSegmentationMode.REMOVE_SMALL_ISLANDS.value:
-            self.effect.apply(
-                self.segmentation_editor.get_active_segment_id(),
-                max_number_of_segments=None,
-                minimum_size=self._typed_state.data.minimum_size,
-                split=False,
-            )
+            self.effect.remove_small_islands(self._typed_state.data.minimum_size)
         elif self._typed_state.data.mode == IslandsSegmentationMode.SPLIT_TO_SEGMENTS.value:
-            self.effect.apply(
-                segment_id=self.segmentation_editor.get_active_segment_id(),
-                max_number_of_segments=None,
-                split=True,
-            )
+            self.effect.split_islands_to_segments()
diff --git a/tests/test_segmentation_islands_effect.py b/tests/test_segmentation_islands_effect.py
@@ -7,13 +7,6 @@
 from trame_slicer.segmentation import SegmentationEffectIslands
 
 
-@pytest.fixture
-def undo_stack(a_segmentation_editor):
-    undo_stack = UndoStack()
-    a_segmentation_editor.set_undo_stack(undo_stack)
-    return undo_stack
-
-
 @pytest.fixture
 def a_segmentation_spheres_file_path(a_data_folder) -> Path:
     return a_data_folder.joinpath("segmentation_spheres.nii.gz")
@@ -35,6 +28,8 @@ def set_up(a_slicer_app, a_volume_node, a_segmentation_editor, a_segmentation_sp
     a_slicer_app.display_manager.show_volume(a_volume_node, vr_preset="MR-Default")
     segmentation_node = a_slicer_app.io_manager.load_segmentation(a_segmentation_spheres_file_path)
     a_segmentation_editor.set_active_segmentation(segmentation_node, a_volume_node)
+    undo_stack = UndoStack()
+    a_segmentation_editor.set_undo_stack(undo_stack)
 
 
 def get_segment_array(
@@ -44,7 +39,7 @@ def get_segment_array(
     return a_segmentation_editor.get_segment_labelmap(segment_id, as_numpy_array=True)
 
 
-def test_do_nothing(
+def test_keep_biggest_island(
     a_segmentation_editor,
     effect,
     segment_id,
@@ -54,60 +49,49 @@ def test_do_nothing(
         a_segmentation_editor,
         segment_id,
     )
-    effect.apply(segment_id, None, minimum_size=0, split=False)
-    assert np.array_equal(
-        get_segment_array(
-            a_segmentation_editor,
-            segment_id,
-        ),
-        source_array,
-    )
-
-
-def test_keep_zero_islands(
-    a_segmentation_editor,
-    effect,
-    segment_id,
-):
-    assert effect.is_active
-    effect.apply(segment_id, 0, minimum_size=0, split=False)
+    effect.keep_largest_island()
     segment_array = get_segment_array(
         a_segmentation_editor,
         segment_id,
     )
-    assert np.array_equal(
-        segment_array,
-        np.zeros_like(segment_array),
-    )
+    # Assert that application created new zeros
+    assert np.count_nonzero(source_array) > np.count_nonzero(segment_array)
 
 
-def test_split(
+def test_split_islands_to_segments(
     a_segmentation_editor,
     effect,
-    segment_id,
 ):
     assert effect.is_active
-    effect.apply(segment_id, None, minimum_size=0, split=True)
+    effect.split_islands_to_segments()
     assert len(a_segmentation_editor.get_segment_ids()) == 3
 
 
-def test_pixel_islands(
+def test_with_0_min_voxel_size_remove_small_islands_does_nothing(
     a_segmentation_editor,
     effect,
     segment_id,
 ):
     assert effect.is_active
-    effect.apply(segment_id, None, minimum_size=1, split=True)
-    assert len(a_segmentation_editor.get_segment_ids()) == 3
+    source_array = get_segment_array(
+        a_segmentation_editor,
+        segment_id,
+    )
+    effect.remove_small_islands(1)
+    segment_array = get_segment_array(
+        a_segmentation_editor,
+        segment_id,
+    )
+    assert np.array_equal(source_array, segment_array)
 
 
-def test_remove_all_islands(
+def test_with_max_min_voxel_size_remove_small_islands_removes_all_islands(
     a_segmentation_editor,
     effect,
     segment_id,
 ):
     assert effect.is_active
-    effect.apply(segment_id, None, minimum_size=int(1e15), split=True)
+    effect.remove_small_islands(int(1e15))
     segment_array = get_segment_array(
         a_segmentation_editor,
         segment_id,
diff --git a/trame_slicer/segmentation/segmentation_effect_islands.py b/trame_slicer/segmentation/segmentation_effect_islands.py
@@ -15,6 +15,7 @@
 
 from . import ModificationMode
 from .segmentation_effect import SegmentationEffect
+from .segmentation_undo_command import SegmentationLabelMapUndoCommand
 
 
 class SegmentationEffectIslands(SegmentationEffect):
@@ -26,62 +27,86 @@ def _create_pipeline(self, _view_node: vtkMRMLAbstractViewNode, _parameter: vtkM
         # Islands effect does not require a pipeline
         return None
 
-    def apply(
-        self, segment_id: str, max_number_of_segments: int | None = None, minimum_size: int = 0, split: bool = True
-    ):
+    def _remove_small_islands(self, minimum_size: int) -> None:
         if not self.is_active:
             return
 
-        island_image = self.get_island_labelmap(segment_id, minimum_size)
+        island_image = self.get_island_labelmap(minimum_size)
+        self.modifier.apply_labelmap(island_image)
 
+    def _keep_n_biggest_islands(self, number_of_islands: int):
+        if not self.is_active:
+            return
+
+        modifier_image = self.modifier.create_modifier_labelmap()
+        if number_of_islands <= 0:
+            self.modifier.apply_labelmap(modifier_image)
+            return
+
+        island_image = self.get_island_labelmap()
         label_values = vtkIntArray()
         vtkSlicerSegmentationsModuleLogic.GetAllLabelValues(label_values, island_image)
         max_label_value = label_values.GetNumberOfTuples()
-        if max_number_of_segments is not None:
-            max_label_value = min(max_label_value, max_number_of_segments)
-        if not split and max_number_of_segments is None:
-            # Overwrite selected segment's labelmap with filtered labelmap
-            self.modifier.apply_labelmap(island_image)
-        else:
-            modifier_image = self.modifier.create_modifier_labelmap()
-            threshold = vtkImageThreshold()
-            threshold.SetInputData(island_image)
-
-            kept_labels = [int(label_values.GetTuple1(i)) for i in range(1, max_label_value)]
-
-            for i in range(max_label_value):
-                label_value = int(label_values.GetTuple1(i))
-                if i == 0:
-                    # Replace selected segment's labelmap by first island
-                    threshold.ThresholdBetween(label_value, label_value)
-                    threshold.SetInValue(1)
-                    threshold.SetOutValue(0)
-                    threshold.Update()
-                    modifier_image.DeepCopy(threshold.GetOutput())
-                # Remove non-kept labels
-                if label_value in kept_labels:
-                    continue
-                threshold.ReplaceOutOff()
-                threshold.ThresholdBetween(label_value, label_value)
-                threshold.SetInValue(0)
-                threshold.Update()
-                threshold.SetInputData(threshold.GetOutput())
+        if number_of_islands >= max_label_value:
+            return
 
-            self.modifier.apply_labelmap(modifier_image)
+        modifier_image = self.modifier.create_modifier_labelmap()
+        threshold = vtkImageThreshold()
+        threshold.SetInputData(island_image)
+        threshold.ReplaceOutOff()
+
+        for i in range(number_of_islands, max_label_value):
+            label_value = int(label_values.GetTuple1(i))
+            threshold.ThresholdBetween(label_value, label_value)
+            threshold.SetInValue(0)
+            threshold.Update()
+            threshold.SetInputData(threshold.GetOutput())
+
+        modifier_image.DeepCopy(threshold.GetOutput())
+        self.modifier.apply_labelmap(modifier_image)
+
+    def _split_islands_to_segments(self):  # TODO: Make it undoable
+        if not self.is_active:
+            return
+
+        island_image = self.get_island_labelmap()
+        label_values = vtkIntArray()
+        vtkSlicerSegmentationsModuleLogic.GetAllLabelValues(label_values, island_image)
+
+        modifier_image = self.modifier.create_modifier_labelmap()
+
+        threshold = vtkImageThreshold()
+        threshold.SetInputData(island_image)
 
-            if max_label_value > 1:
+        # Replace selected segment's labelmap by first island
+        initial_label_value = int(label_values.GetTuple1(0))
+        threshold.ThresholdBetween(initial_label_value, initial_label_value)
+        threshold.SetInValue(initial_label_value)
+        threshold.SetOutValue(0)
+        threshold.Update()
+        modifier_image.DeepCopy(threshold.GetOutput())
+
+        # Create labelmap without first segment
+        threshold.SetInValue(0)
+        threshold.ReplaceOutOff()
+        threshold.ThresholdBetween(initial_label_value, initial_label_value)
+        threshold.Update()
+
+        with self.modifier.group_undo_commands(f"{__class__} - Split {self.modifier.active_segment_id}"):
+            self.modifier.apply_labelmap(modifier_image)
+            with SegmentationLabelMapUndoCommand.push_state_change(self.modifier.segmentation):
                 modifier_image.DeepCopy(threshold.GetOutput())
                 vtkSlicerSegmentationsModuleLogic.ImportLabelmapToSegmentationNode(
                     modifier_image,
                     self.modifier.segmentation.segmentation_node,
-                    self.modifier.segmentation.get_segment(segment_id).GetName(),
+                    self.modifier.segmentation.get_segment(self.modifier.active_segment_id).GetName(),
                 )
-            self.modifier.segmentation.segmentation_modified.emit()
+        self.modifier.segmentation.segmentation_modified.emit()
 
-    def get_island_labelmap(self, segment_id: str, minimum_size: int = 0) -> vtkOrientedImageData:
+    def get_island_labelmap(self, minimum_size: int = 0) -> vtkOrientedImageData:
         source_image_data = self.modifier.get_source_image_data()
 
-        segment_labelmap = self.modifier.get_segment_labelmap(segment_id)
+        segment_labelmap = self.modifier.get_segment_labelmap(self.modifier.active_segment_id)
         cast_in = vtkImageCast()
         cast_in.SetInputData(segment_labelmap)
         cast_in.SetOutputScalarTypeToUnsignedInt()
@@ -100,3 +125,12 @@ def get_island_labelmap(self, segment_id: str, minimum_size: int = 0) -> vtkOrie
         island_image.SetImageToWorldMatrix(image_to_world_matrix)
 
         return island_image
+
+    def keep_largest_island(self) -> None:
+        self._keep_n_biggest_islands(1)  # TODO: test, pass to 1
+
+    def remove_small_islands(self, min_voxel_size: int) -> None:
+        self._remove_small_islands(min_voxel_size)
+
+    def split_islands_to_segments(self) -> None:
+        self._split_islands_to_segments()
