feat(segmentation): add draw segmentation effect

Author: Jo-Byr

examples/viewer_lib/ui/segmentation/segment_editor_ui.py

@@ -17,6 +17,7 @@
 
 from trame_slicer.segmentation import (
     SegmentationEffect,
+    SegmentationEffectDraw,
     SegmentationEffectErase,
     SegmentationEffectIslands,
     SegmentationEffectNoTool,
@@ -138,6 +139,12 @@ def build_effect_buttons(self, all: bool = True, **kwargs):
             SegmentationEffectScissors,
             **kwargs,
         )
+        self._create_effect_button(
+            "Draw",
+            "mdi-pencil",
+            SegmentationEffectDraw,
+            **kwargs,
+        )
         if all:
             self._create_effect_button(
                 "Threshold",

tests/test_segmentation_draw_effect.py

@@ -0,0 +1,46 @@
+import pytest
+
+from tests.conftest import a_slice_view, a_threed_view
+from tests.view_events import MouseButton, ViewEvents
+from trame_slicer.segmentation import SegmentationEffectDraw
+
+
+def apply_draw_effect(view):
+    view_events = ViewEvents(view)
+    center_x, center_y = view_events.view_center()
+    view_events.mouse_move_to(center_x, center_y)
+    view_events.mouse_press_event()
+    view_events.mouse_release_event()
+    view_events.mouse_move_to(0, center_y)
+    view_events.mouse_press_event()
+    view_events.mouse_release_event()
+    view_events.mouse_move_to(0, 0)
+    view_events.mouse_press_event()
+    view_events.mouse_release_event()
+    view_events.mouse_press_event(MouseButton.Right)
+
+
+@pytest.mark.parametrize("view", [a_slice_view, a_threed_view])
+def test_draw_effect_adds_segmentation_to_selected_segment(
+    a_slicer_app,
+    a_segmentation_editor,
+    a_volume_node,
+    view,
+    request,
+    render_interactive,
+):
+    view = request.getfixturevalue(view.__name__)
+    a_slicer_app.display_manager.show_volume(a_volume_node, vr_preset="MR-Default")
+
+    segmentation_node = a_segmentation_editor.create_empty_segmentation_node()
+    a_segmentation_editor.set_active_segmentation(segmentation_node, a_volume_node)
+    segment_id = a_segmentation_editor.add_empty_segment()
+    a_segmentation_editor.set_active_segment_id(segment_id)
+    a_segmentation_editor.set_active_effect_type(SegmentationEffectDraw)
+    apply_draw_effect(view)
+    array = a_segmentation_editor.get_segment_labelmap(segment_id, as_numpy_array=True)
+    assert array.sum() > 0
+
+    if render_interactive:
+        a_segmentation_editor.show_3d(True)
+        view.interactor().Start()

trame_slicer/core/segmentation_editor.py

@@ -31,6 +31,7 @@
     Segmentation,
     SegmentationDisplay,
     SegmentationEffect,
+    SegmentationEffectDraw,
     SegmentationEffectErase,
     SegmentationEffectIslands,
     SegmentationEffectNoTool,
@@ -54,6 +55,7 @@ class SegmentationEditor(SignalContainer):
     """
 
     builtin_effects: ClassVar[list[type[SegmentationEffect]]] = [
+        SegmentationEffectDraw,
         SegmentationEffectErase,
         SegmentationEffectIslands,
         SegmentationEffectNoTool,

trame_slicer/segmentation/__init__.py

@@ -2,11 +2,14 @@
 
 from .brush_source import BrushSource
 from .paint_effect_parameters import BrushDiameterMode, BrushShape
+from .polygon_brush import PolygonBrush
 from .segment_modifier import ModificationMode, SegmentModifier
 from .segment_properties import SegmentProperties
 from .segmentation import Segmentation
 from .segmentation_display import SegmentationDisplay, SegmentationOpacityEnum
 from .segmentation_effect import SegmentationEffect
+from .segmentation_effect_draw import SegmentationEffectDraw
+from .segmentation_effect_draw_widget import SegmentationDrawPipeline
 from .segmentation_effect_islands import SegmentationEffectIslands
 from .segmentation_effect_no_tool import SegmentationEffectNoTool
 from .segmentation_effect_paint_erase import (
@@ -16,11 +19,7 @@
 )
 from .segmentation_effect_pipeline import SegmentationEffectPipeline
 from .segmentation_effect_scissors import SegmentationEffectScissors
-from .segmentation_effect_scissors_widget import (
-    ScissorsPolygonBrush,
-    SegmentationScissorsPipeline,
-    SegmentationScissorsWidget,
-)
+from .segmentation_effect_scissors_widget import SegmentationScissorsPipeline
 from .segmentation_effect_threshold import (
     AutoThresholdMethod,
     AutoThresholdMode,
@@ -37,6 +36,10 @@
     SegmentationPaintWidget2D,
     SegmentationPaintWidget3D,
 )
+from .segmentation_polygon_widget import (
+    SegmentationPolygonPipeline,
+    SegmentationPolygonWidget,
+)
 
 __all__ = [
     "AutoThresholdMethod",
@@ -45,12 +48,14 @@
     "BrushShape",
     "BrushSource",
     "ModificationMode",
-    "ScissorsPolygonBrush",
+    "PolygonBrush",
     "SegmentModifier",
     "SegmentProperties",
     "Segmentation",
     "SegmentationDisplay",
+    "SegmentationDrawPipeline",
     "SegmentationEffect",
+    "SegmentationEffectDraw",
     "SegmentationEffectErase",
     "SegmentationEffectIslands",
     "SegmentationEffectNoTool",
@@ -65,8 +70,9 @@
     "SegmentationPaintWidget",
     "SegmentationPaintWidget2D",
     "SegmentationPaintWidget3D",
+    "SegmentationPolygonPipeline",
+    "SegmentationPolygonWidget",
     "SegmentationScissorsPipeline",
-    "SegmentationScissorsWidget",
     "SegmentationThresholdPipeline2D",
     "ThresholdParameters",
 ]

trame_slicer/segmentation/polygon_brush.py

@@ -0,0 +1,71 @@
+from vtkmodules.vtkCommonCore import vtkPoints
+from vtkmodules.vtkCommonDataModel import vtkCellArray, vtkPolyData
+from vtkmodules.vtkRenderingCore import (
+    vtkActor2D,
+    vtkPolyDataMapper2D,
+    vtkProp,
+    vtkProperty2D,
+)
+
+
+class PolygonBrush:
+    """Display the draw polygon as 2D lines"""
+
+    def __init__(self):
+        super().__init__()
+        self._points = vtkPoints()
+        self._lines = vtkCellArray()
+        self._vertices = vtkCellArray()
+        self._poly = vtkPolyData()
+        self._poly.SetLines(self._lines)
+        self._poly.SetVerts(self._vertices)
+        self._poly.SetPoints(self._points)
+
+        self._brush_mapper = vtkPolyDataMapper2D()
+        self._brush_mapper.SetInputData(self._poly)
+        self._brush_actor = vtkActor2D()
+        self._brush_actor.SetMapper(self._brush_mapper)
+        self._brush_actor.VisibilityOff()
+        props = self._brush_actor.GetProperty()
+        props.SetColor(1.0, 1.0, 0.0)
+        props.SetPointSize(4.0)
+        props.SetLineWidth(2.0)
+
+    def set_visibility(self, visible: bool):
+        self._brush_actor.SetVisibility(int(visible))
+
+    def move_last_point(self, x: int, y: int) -> None:
+        count = self._points.GetNumberOfPoints()
+        if count == 0:
+            self.add_point(x, y)
+        else:
+            self._points.SetPoint(count - 1, [float(x), float(y), 1.0])
+            self._points.Modified()
+
+    def add_point(self, x: int, y: int) -> None:
+        self._points.InsertNextPoint([float(x), float(y), 1.0])
+        count = self._points.GetNumberOfPoints()
+        if count > 1:
+            self._lines.InsertNextCell(2, [count - 1, count - 2])
+        self._vertices.InsertNextCell(1, [count - 1])
+        self._points.Modified()
+
+    def reset(self) -> None:
+        self._points.SetNumberOfPoints(0)
+        self._lines.Reset()
+        self._vertices.Reset()
+        self._poly.Modified()
+
+    @property
+    def points(self) -> vtkPoints:
+        return self._points
+
+    def get_prop(self) -> vtkProp:
+        """
+        Return brush prop.
+        Can be used to add or remove the brush from the renderer, configure rendering properties (visibility, color, ...)
+        """
+        return self._brush_actor
+
+    def get_property(self) -> vtkProperty2D:
+        return self._brush_actor.GetProperty()

trame_slicer/segmentation/segmentation_effect_draw.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from slicer import vtkMRMLAbstractViewNode, vtkMRMLNode
+
+from .segmentation_effect import SegmentationEffect
+from .segmentation_effect_draw_widget import SegmentationDrawPipeline
+from .segmentation_effect_pipeline import SegmentationEffectPipeline
+
+
+class SegmentationEffectDraw(SegmentationEffect):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def _create_pipeline(
+        self, _view_node: vtkMRMLAbstractViewNode, _parameter: vtkMRMLNode
+    ) -> SegmentationEffectPipeline | None:
+        return SegmentationDrawPipeline()

trame_slicer/segmentation/segmentation_effect_draw_widget.py

@@ -0,0 +1,37 @@
+from slicer import vtkMRMLInteractionEventData
+
+from .segmentation_polygon_widget import SegmentationPolygonPipeline
+
+
+class SegmentationDrawPipeline(SegmentationPolygonPipeline):
+    def _LeftButtonPressed(self, event_data: vtkMRMLInteractionEventData) -> bool:
+        x, y = event_data.GetDisplayPosition()
+        if self.widget.is_painting():
+            self.widget.add_point(x, y)
+        else:
+            self.widget.start_painting(x, y)
+        self.RequestRender()
+        return True
+
+    def _LeftButtonReleased(self, _event_data: vtkMRMLInteractionEventData) -> bool:
+        self.widget.pause_painting()
+        return True
+
+    def _MouseMoved(self, event_data: vtkMRMLInteractionEventData) -> bool:
+        if self.widget.is_painting():
+            x, y = event_data.GetDisplayPosition()
+            self.widget.move_last_point(x, y)
+            if self.widget.is_painting():
+                self.widget.add_point(x, y)
+            self.RequestRender()
+
+        # Always let other interactor and displayable managers do whatever they want
+        return False
+
+    def _RightButtonPressed(self, _event_data: vtkMRMLInteractionEventData) -> bool:
+        self.widget.stop_painting()
+        self.RequestRender()
+        return True
+
+    def _RightButtonReleased(self, _event_data: vtkMRMLInteractionEventData) -> bool:
+        return True