feat(segmentation): add island segmentation effect

examples/viewer_lib/logic/__init__.py

@@ -4,6 +4,7 @@
 from .medical_viewer_logic import MedicalViewerLogic
 from .segmentation import (
     EraseEffectLogic,
+    IslandsEffectLogic,
     PaintEffectLogic,
     PaintEraseEffectLogic,
     SegmentEditDialogLogic,
@@ -15,6 +16,7 @@
 __all__ = [
     "BaseLogic",
     "EraseEffectLogic",
+    "IslandsEffectLogic",
     "LoadFilesLogic",
     "MarkupsButtonLogic",
     "MedicalViewerLogic",

examples/viewer_lib/logic/segmentation/__init__.py

@@ -1,4 +1,5 @@
 from .base_segmentation_logic import BaseEffectLogic, BaseSegmentationLogic
+from .islands_effect_logic import IslandsEffectLogic
 from .paint_erase_effect_logic import (
     EraseEffectLogic,
     PaintEffectLogic,
@@ -12,6 +13,7 @@
     "BaseEffectLogic",
     "BaseSegmentationLogic",
     "EraseEffectLogic",
+    "IslandsEffectLogic",
     "PaintEffectLogic",
     "PaintEraseEffectLogic",
     "SegmentEditDialogLogic",

examples/viewer_lib/logic/segmentation/islands_effect_logic.py

@@ -0,0 +1,33 @@
+from trame_server import Server
+
+from trame_slicer.core import SlicerApp
+from trame_slicer.segmentation import SegmentationEffectIslands
+
+from ...ui import (
+    IslandsEffectUI,
+    IslandsSegmentationMode,
+    IslandsState,
+    SegmentEditorUI,
+)
+from .base_segmentation_logic import BaseEffectLogic
+
+
+class IslandsEffectLogic(BaseEffectLogic[IslandsState, SegmentationEffectIslands]):
+    def __init__(self, server: Server, slicer_app: SlicerApp):
+        super().__init__(server, slicer_app, IslandsState, SegmentationEffectIslands)
+
+    def set_ui(self, ui: SegmentEditorUI):
+        self.set_effect_ui(ui.get_effect_ui(SegmentationEffectIslands))
+
+    def set_effect_ui(self, islands_ui: IslandsEffectUI):
+        islands_ui.apply_clicked.connect(self._on_apply_clicked)
+
+    def _on_apply_clicked(self):
+        if not self.is_active():
+            return
+        if self._typed_state.data.mode == IslandsSegmentationMode.KEEP_LARGEST_ISLAND:
+            self.effect.keep_largest_island()
+        elif self._typed_state.data.mode == IslandsSegmentationMode.REMOVE_SMALL_ISLANDS:
+            self.effect.remove_small_islands(self._typed_state.data.minimum_size)
+        elif self._typed_state.data.mode == IslandsSegmentationMode.SPLIT_TO_SEGMENTS:
+            self.effect.split_islands_to_segments()

examples/viewer_lib/logic/segmentation/segment_editor_logic.py

@@ -17,6 +17,7 @@
     get_current_volume_node,
 )
 from .base_segmentation_logic import BaseEffectLogic, BaseSegmentationLogic
+from .islands_effect_logic import IslandsEffectLogic
 from .paint_erase_effect_logic import EraseEffectLogic, PaintEffectLogic
 from .segment_edit_dialog_logic import SegmentEditDialogLogic
 from .threshold_effect_logic import ThresholdEffectLogic
@@ -27,6 +28,7 @@ def __init__(self, server: Server, slicer_app: SlicerApp):
         super().__init__(server=server, slicer_app=slicer_app, state_type=SegmentEditorState)
 
         effect_logic = [
+            IslandsEffectLogic,
             ThresholdEffectLogic,
             PaintEffectLogic,
             EraseEffectLogic,

examples/viewer_lib/ui/__init__.py

@@ -10,6 +10,9 @@
 from .medical_viewer_ui import MedicalViewerUI
 from .mpr_interaction_button import MprInteractionButton, MprInteractionButtonState
 from .segmentation import (
+    IslandsEffectUI,
+    IslandsSegmentationMode,
+    IslandsState,
     PaintEffectState,
     PaintEffectUI,
     SegmentationOpacityUI,
@@ -35,6 +38,9 @@
 __all__ = [
     "ControlButton",
     "IdName",
+    "IslandsEffectUI",
+    "IslandsSegmentationMode",
+    "IslandsState",
     "LayoutButton",
     "LayoutButton",
     "LayoutButtonState",

examples/viewer_lib/ui/segmentation/__init__.py

@@ -1,3 +1,4 @@
+from .islands_effect_ui import IslandsEffectUI, IslandsSegmentationMode, IslandsState
 from .paint_effect_ui import PaintEffectState, PaintEffectUI
 from .segment_edit_dialog import SegmentEditDialog, SegmentEditDialogState
 from .segment_editor_ui import SegmentEditorState, SegmentEditorUI
@@ -7,6 +8,9 @@
 from .threshold_effect_ui import ThresholdEffectUI, ThresholdState
 
 __all__ = [
+    "IslandsEffectUI",
+    "IslandsSegmentationMode",
+    "IslandsState",
     "PaintEffectState",
     "PaintEffectUI",
     "SegmentEditDialog",

examples/viewer_lib/ui/segmentation/islands_effect_ui.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+from enum import Enum, auto
+
+from trame_server.utils.typed_state import TypedState
+from trame_vuetify.widgets.vuetify3 import VBtn, VContainer, VNumberInput, VSelect
+from undo_stack import Signal
+
+
+class IslandsSegmentationMode(Enum):
+    KEEP_LARGEST_ISLAND = auto()
+    REMOVE_SMALL_ISLANDS = auto()
+    SPLIT_TO_SEGMENTS = auto()
+
+
+@dataclass
+class IslandsState:
+    mode: IslandsSegmentationMode = IslandsSegmentationMode.KEEP_LARGEST_ISLAND
+    minimum_size: int = 1000
+
+
+class IslandsEffectUI(VContainer):
+    apply_clicked = Signal()
+
+    def __init__(self, **kwargs):
+        super().__init__(classes="fill-width", **kwargs)
+        self._typed_state = TypedState(self.state, IslandsState)
+
+        self.labels = {
+            IslandsSegmentationMode.KEEP_LARGEST_ISLAND: "Keep largest island",
+            IslandsSegmentationMode.REMOVE_SMALL_ISLANDS: "Remove small islands",
+            IslandsSegmentationMode.SPLIT_TO_SEGMENTS: "Split to segments",
+        }
+
+        with self:
+            VSelect(
+                v_model=self._typed_state.name.mode,
+                items=(
+                    [
+                        {"title": self.labels[mode], "value": self._typed_state.encode(mode)}
+                        for mode in IslandsSegmentationMode
+                    ],
+                ),
+                hide_details=True,
+                density="compact",
+                label="Mode",
+            )
+            VNumberInput(
+                v_model=self._typed_state.name.minimum_size,
+                label="Minimum size",
+                disabled=(
+                    f"{self._typed_state.name.mode} !== {self._typed_state.encode(IslandsSegmentationMode.REMOVE_SMALL_ISLANDS)}",
+                ),
+                min=(0,),
+                hide_details=True,
+                density="compact",
+                classes="mt-5",
+            )
+            VBtn("Apply", prepend_icon="mdi-check-outline", block=True, click=self.apply_clicked, classes="mt-5")

examples/viewer_lib/ui/segmentation/segment_editor_ui.py

@@ -15,13 +15,15 @@
 from trame_slicer.segmentation import (
     SegmentationEffect,
     SegmentationEffectErase,
+    SegmentationEffectIslands,
     SegmentationEffectNoTool,
     SegmentationEffectPaint,
     SegmentationEffectScissors,
     SegmentationEffectThreshold,
 )
 
 from ..control_button import ControlButton
+from .islands_effect_ui import IslandsEffectUI
 from .paint_effect_ui import PaintEffectUI
 from .segment_edit_dialog import SegmentEditDialog, SegmentEditDialogState
 from .segment_list import SegmentList, SegmentListState
@@ -105,12 +107,14 @@ def __init__(self, **kwargs):
                 self._create_effect_button("Erase", "mdi-eraser", SegmentationEffectErase)
                 self._create_effect_button("Scissors", "mdi-content-cut", SegmentationEffectScissors)
                 self._create_effect_button("Threshold", "mdi-auto-fix", SegmentationEffectThreshold)
+                self._create_effect_button("Islands", "mdi-scatter-plot", SegmentationEffectIslands)
                 VDivider()
 
             with VRow():
                 self._register_effect_ui(SegmentationEffectThreshold, ThresholdEffectUI)
                 self._register_effect_ui(SegmentationEffectPaint, PaintEffectUI)
                 self._register_effect_ui(SegmentationEffectErase, PaintEffectUI)
+                self._register_effect_ui(SegmentationEffectIslands, IslandsEffectUI)
 
             with VRow():
                 VBtn(

tests/examples/test_segment_islands.py

@@ -0,0 +1,38 @@
+import pytest
+
+from examples.viewer_lib.logic import IslandsEffectLogic
+from examples.viewer_lib.ui import (
+    IslandsEffectUI,
+    IslandsSegmentationMode,
+    MedicalViewerLayout,
+)
+
+
+@pytest.fixture
+def effect_ui(a_server):
+    with MedicalViewerLayout(a_server, is_drawer_visible=True) as ui, ui.drawer:
+        return IslandsEffectUI()
+
+
+@pytest.fixture
+def effect_logic(a_server, a_slicer_app, effect_ui):
+    logic = IslandsEffectLogic(a_server, a_slicer_app)
+    logic.set_effect_ui(effect_ui)
+    return logic
+
+
+@pytest.mark.parametrize("island_mode", list(IslandsSegmentationMode))
+def test_can_apply_island_effect(
+    effect_logic,
+    effect_ui,
+    a_segmentation_nifti_file_path,
+    a_segmentation_editor,
+    a_slicer_app,
+    a_volume_node,
+    island_mode,
+):
+    segmentation_node = a_slicer_app.io_manager.load_segmentation(a_segmentation_nifti_file_path)
+    a_segmentation_editor.set_active_segmentation(segmentation_node, a_volume_node)
+    effect_logic.set_active()
+    effect_ui._typed_state.data.mode = island_mode
+    effect_ui.apply_clicked()

tests/test_segmentation_islands_effect.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+from undo_stack import UndoStack
+
+from trame_slicer.segmentation import SegmentationEffectIslands
+
+
+@pytest.fixture
+def a_segmentation_spheres_file_path(a_data_folder) -> Path:
+    return a_data_folder.joinpath("segmentation_spheres.nii.gz")
+
+
+@pytest.fixture
+def effect(a_segmentation_editor):
+    effect: SegmentationEffectIslands = a_segmentation_editor.set_active_effect_type(SegmentationEffectIslands)
+    return effect
+
+
+@pytest.fixture
+def segment_id(a_segmentation_editor):
+    return a_segmentation_editor.get_nth_segment_id(0)
+
+
+@pytest.fixture(autouse=True)
+def set_up(a_slicer_app, a_volume_node, a_segmentation_editor, a_segmentation_spheres_file_path):
+    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(a_segmentation_editor, segment_id):
+    return a_segmentation_editor.get_segment_labelmap(segment_id, as_numpy_array=True)
+
+
+def test_keep_biggest_island(a_segmentation_editor, effect, segment_id):
+    assert effect.is_active
+    source_array = get_segment_array(a_segmentation_editor, segment_id)
+    effect.keep_largest_island()
+    segment_array = get_segment_array(a_segmentation_editor, segment_id)
+    # Assert that application created new zeros
+    assert np.count_nonzero(source_array) > np.count_nonzero(segment_array)
+
+
+def test_split_islands_to_segments(a_segmentation_editor, effect):
+    assert effect.is_active
+    effect.split_islands_to_segments()
+    assert len(a_segmentation_editor.get_segment_ids()) == 3
+
+
+def test_with_0_min_voxel_size_remove_small_islands_does_nothing(a_segmentation_editor, effect, segment_id):
+    assert effect.is_active
+    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_with_max_min_voxel_size_remove_small_islands_removes_all_islands(a_segmentation_editor, effect, segment_id):
+    assert effect.is_active
+    effect.remove_small_islands(int(1e15))
+    segment_array = get_segment_array(a_segmentation_editor, segment_id)
+    assert np.array_equal(segment_array, np.zeros_like(segment_array))

trame_slicer/core/segmentation_editor.py

@@ -32,6 +32,7 @@
     SegmentationDisplay,
     SegmentationEffect,
     SegmentationEffectErase,
+    SegmentationEffectIslands,
     SegmentationEffectNoTool,
     SegmentationEffectPaint,
     SegmentationEffectPipeline,
@@ -54,6 +55,7 @@ class SegmentationEditor(SignalContainer):
 
     builtin_effects: ClassVar[list[type[SegmentationEffect]]] = [
         SegmentationEffectErase,
+        SegmentationEffectIslands,
         SegmentationEffectNoTool,
         SegmentationEffectPaint,
         SegmentationEffectScissors,

trame_slicer/segmentation/__init__.py

@@ -7,6 +7,7 @@
 from .segmentation import Segmentation
 from .segmentation_display import SegmentationDisplay, SegmentationOpacityEnum
 from .segmentation_effect import SegmentationEffect
+from .segmentation_effect_islands import SegmentationEffectIslands
 from .segmentation_effect_no_tool import SegmentationEffectNoTool
 from .segmentation_effect_paint_erase import (
     SegmentationEffectErase,
@@ -51,6 +52,7 @@
     "SegmentationDisplay",
     "SegmentationEffect",
     "SegmentationEffectErase",
+    "SegmentationEffectIslands",
     "SegmentationEffectNoTool",
     "SegmentationEffectPaint",
     "SegmentationEffectPaintErase",

trame_slicer/segmentation/segment_modifier.py

@@ -3,6 +3,7 @@
 import enum
 import logging
 import math
+from collections.abc import Generator
 from enum import auto
 
 from numpy.typing import NDArray
@@ -380,3 +381,6 @@ def is_source_volume_intensity_mask_enabled(self) -> bool:
         if not self.segment_editor_node:
             return False
         return self.segment_editor_node.GetSourceVolumeIntensityMask()
+
+    def group_undo_commands(self, text: str = "") -> Generator:
+        return self.segmentation.group_undo_commands(text)

trame_slicer/segmentation/segmentation.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Generator
+from contextlib import contextmanager
+
 from numpy.typing import NDArray
 from slicer import (
     vtkMRMLSegmentationNode,
@@ -272,3 +275,12 @@ def set_segment_labelmap(self, segment_id, label_map: vtkImageData | NDArray):
 
     def get_display(self) -> SegmentationDisplay | None:
         return SegmentationDisplay(self._segmentation_node.GetDisplayNode()) if self._segmentation_node else None
+
+    @contextmanager
+    def group_undo_commands(self, text: str = "") -> Generator:
+        if not self.undo_stack:
+            yield
+            return
+
+        with self.undo_stack.group_undo_commands(text):
+            yield