feat(segmentation): Enhance Segmentation with New AI and Once Click Tools (#4910)

2025-04-01 14:19:03 -04:00
1 parent 908482d2e0
commit c75a3f3af5
@@ -38,7 +38,7 @@
    "@cornerstonejs/codec-libjpeg-turbo-8bit": "^1.2.2",
    "@cornerstonejs/codec-openjpeg": "^1.2.4",
    "@cornerstonejs/codec-openjph": "^2.4.5",
-    "@cornerstonejs/dicom-image-loader": "^3.7.17",
+    "@cornerstonejs/dicom-image-loader": "^3.8.2",
    "@ohif/core": "3.10.0-beta.143",
    "@ohif/ui": "3.10.0-beta.143",
    "dcmjs": "*",
@@ -51,11 +51,12 @@
  },
  "dependencies": {
    "@babel/runtime": "^7.20.13",
-    "@cornerstonejs/adapters": "^3.7.17",
-    "@cornerstonejs/core": "^3.7.17",
-    "@cornerstonejs/labelmap-interpolation": "^3.7.17",
-    "@cornerstonejs/polymorphic-segmentation": "^3.7.17",
-    "@cornerstonejs/tools": "^3.7.17",
+    "@cornerstonejs/adapters": "^3.8.2",
+    "@cornerstonejs/core": "^3.8.2",
+    "@cornerstonejs/labelmap-interpolation": "^3.8.2",
+    "@cornerstonejs/ai": "^3.8.2",
+    "@cornerstonejs/polymorphic-segmentation": "^3.8.2",
+    "@cornerstonejs/tools": "^3.8.2",
    "@itk-wasm/morphological-contour-interpolation": "1.1.0",
    "@kitware/vtk.js": "32.12.0",
    "html2canvas": "^1.4.1",
@@ -16,6 +16,7 @@ import {
 } from '@cornerstonejs/tools';
 import * as cornerstoneTools from '@cornerstonejs/tools';
 import * as labelmapInterpolation from '@cornerstonejs/labelmap-interpolation';
+import { ONNXSegmentationController } from '@cornerstonejs/ai';

 import { Types as OhifTypes, utils } from '@ohif/core';
 import i18n from '@ohif/i18n';
@@ -59,6 +60,28 @@ const getLabelmapTools = ({ toolGroupService }) => {
  return labelmapTools;
 };

+const segmentAI = new ONNXSegmentationController({
+  autoSegmentMode: true,
+  models: {
+    sam_b: [
+      {
+        name: 'sam-b-encoder',
+        url: 'https://huggingface.co/schmuell/sam-b-fp16/resolve/main/sam_vit_b_01ec64.encoder-fp16.onnx',
+        size: 180,
+        key: 'encoder',
+      },
+      {
+        name: 'sam-b-decoder',
+        url: 'https://huggingface.co/schmuell/sam-b-fp16/resolve/main/sam_vit_b_01ec64.decoder.onnx',
+        size: 17,
+        key: 'decoder',
+      },
+    ],
+  },
+  modelName: 'sam_b',
+});
+let segmentAIEnabled = false;
+
 function commandsModule({
  servicesManager,
  commandsManager,
@@ -979,7 +1002,7 @@ function commandsModule({
        });
      }
    },
-    setSourceViewportForReferenceLinesTool: ({ viewportId }) => {
+    setViewportForToolConfiguration: ({ viewportId, toolName }) => {
      if (!viewportId) {
        const { activeViewportId } = viewportGridService.getState();
        viewportId = activeViewportId ?? 'default';
@@ -987,9 +1010,11 @@ function commandsModule({

      const toolGroup = toolGroupService.getToolGroupForViewport(viewportId);

+      const prevConfig = toolGroup?.getToolConfiguration(toolName);
      toolGroup?.setToolConfiguration(
-        ReferenceLinesTool.toolName,
+        toolName,
        {
+          ...prevConfig,
          sourceViewportId: viewportId,
        },
        true // overwrite
@@ -1525,17 +1550,85 @@ function commandsModule({
        }
      });
    },
-    acceptPreview: () => {
+    toggleUseCenterSegmentIndex: ({ toggle }) => {
      const labelmapTools = getLabelmapTools({ toolGroupService });
      labelmapTools.forEach(tool => {
-        tool.acceptPreview();
+        tool.configuration = {
+          ...tool.configuration,
+          useCenterSegmentIndex: toggle,
+        };
      });
    },
-    rejectPreview: () => {
+    _handlePreviewAction: action => {
      const labelmapTools = getLabelmapTools({ toolGroupService });
-      labelmapTools.forEach(tool => {
-        tool.rejectPreview();
+      const { viewport } = _getActiveViewportEnabledElement();
+      const activeTools = labelmapTools.filter(
+        tool => tool.mode === 'Active' || tool.mode === 'Enabled'
+      );
+
+      activeTools.forEach(tool => {
+        tool[`${action}Preview`]();
      });
+
+      if (segmentAI.enabled) {
+        segmentAI[`${action}Preview`](viewport.element);
+      }
+    },
+    acceptPreview: () => {
+      actions._handlePreviewAction('accept');
+    },
+    rejectPreview: () => {
+      actions._handlePreviewAction('reject');
+    },
+    clearMarkersForMarkerLabelmap: () => {
+      const { viewport } = _getActiveViewportEnabledElement();
+      const toolGroup = cornerstoneTools.ToolGroupManager.getToolGroupForViewport(viewport.id);
+      const toolInstance = toolGroup.getToolInstance('MarkerLabelmap');
+
+      if (!toolInstance) {
+        return;
+      }
+
+      toolInstance.clearMarkers(viewport);
+    },
+    interpolateScrollForMarkerLabelmap: () => {
+      const { viewport } = _getActiveViewportEnabledElement();
+      const toolGroup = cornerstoneTools.ToolGroupManager.getToolGroupForViewport(viewport.id);
+      const toolInstance = toolGroup.getToolInstance('MarkerLabelmap');
+
+      if (!toolInstance) {
+        return;
+      }
+
+      toolInstance.interpolateScroll(viewport, 1);
+    },
+    toggleLabelmapAssist: async () => {
+      const { viewport } = _getActiveViewportEnabledElement();
+      const newState = !segmentAI.enabled;
+      segmentAI.enabled = newState;
+
+      if (!segmentAIEnabled) {
+        await segmentAI.initModel();
+        segmentAIEnabled = true;
+      }
+
+      // set the brush tool to active
+      const toolGroupIds = toolGroupService.getToolGroupIds();
+      if (newState) {
+        actions.setToolActiveToolbar({
+          toolName: 'CircularBrushForAutoSegmentAI',
+          toolGroupIds: toolGroupIds,
+        });
+      } else {
+        toolGroupIds.forEach(toolGroupId => {
+          const toolGroup = cornerstoneTools.ToolGroupManager.getToolGroup(toolGroupId);
+          toolGroup.setToolPassive('CircularBrushForAutoSegmentAI');
+        });
+      }
+
+      if (segmentAI.enabled) {
+        segmentAI.initViewport(viewport);
+      }
    },
  };

@@ -1680,8 +1773,8 @@ function commandsModule({
    setViewportColormap: {
      commandFn: actions.setViewportColormap,
    },
-    setSourceViewportForReferenceLinesTool: {
-      commandFn: actions.setSourceViewportForReferenceLinesTool,
+    setViewportForToolConfiguration: {
+      commandFn: actions.setViewportForToolConfiguration,
    },
    storePresentation: {
      commandFn: actions.storePresentation,
@@ -1806,6 +1899,10 @@ function commandsModule({
    toggleSegmentSelect: actions.toggleSegmentSelect,
    acceptPreview: actions.acceptPreview,
    rejectPreview: actions.rejectPreview,
+    toggleUseCenterSegmentIndex: actions.toggleUseCenterSegmentIndex,
+    toggleLabelmapAssist: actions.toggleLabelmapAssist,
+    interpolateScrollForMarkerLabelmap: actions.interpolateScrollForMarkerLabelmap,
+    clearMarkersForMarkerLabelmap: actions.clearMarkersForMarkerLabelmap,
  };

  return {
@@ -36,7 +36,7 @@ export default function getSegmentationPanelCustomization({ commandsManager, ser
      // Get initial states based on current configuration
      const [previewEdits, setPreviewEdits] = useState(false);
      const [toggleSegmentEnabled, setToggleSegmentEnabled] = useState(false);
-
+      const [useCenterAsSegmentIndex, setUseCenterAsSegmentIndex] = useState(false);
      const handlePreviewEditsChange = checked => {
        setPreviewEdits(checked);
        commandsManager.run('toggleSegmentPreviewEdit', { toggle: checked });
@@ -47,6 +47,11 @@ export default function getSegmentationPanelCustomization({ commandsManager, ser
        commandsManager.run('toggleSegmentSelect', { toggle: checked });
      };

+      const handleUseCenterAsSegmentIndexChange = checked => {
+        setUseCenterAsSegmentIndex(checked);
+        commandsManager.run('toggleUseCenterSegmentIndex', { toggle: checked });
+      };
+
      return (
        <div className="bg-muted flex flex-col gap-4 border-b border-b-[2px] border-black px-2 py-3">
          <div className="flex items-center gap-2">
@@ -57,12 +62,20 @@ export default function getSegmentationPanelCustomization({ commandsManager, ser
            <span className="text-base text-white">Preview edits before creating</span>
          </div>

+          <div className="flex items-center gap-2">
+            <Switch
+              checked={useCenterAsSegmentIndex}
+              onCheckedChange={handleUseCenterAsSegmentIndexChange}
+            />
+            <span className="text-base text-white">Use Center as Segment Index</span>
+          </div>
+
          <div className="flex items-center gap-2">
            <Switch
              checked={toggleSegmentEnabled}
              onCheckedChange={handleToggleSegmentEnabledChange}
            />
-            <span className="text-base text-white">Highlight segments to select</span>
+            <span className="text-base text-white">Hover on segment border to activate</span>
          </div>
        </div>
      );
@@ -38,7 +38,9 @@ import {
  OrientationMarkerTool,
  WindowLevelRegionTool,
  SegmentSelectTool,
+  RegionSegmentPlusTool,
 } from '@cornerstonejs/tools';
+import { LabelmapSlicePropagationTool, MarkerLabelmapTool } from '@cornerstonejs/ai';
 import * as polySeg from '@cornerstonejs/polymorphic-segmentation';

 import CalibrationLineTool from './tools/CalibrationLineTool';
@@ -98,6 +100,9 @@ export default function initCornerstoneTools(configuration = {}) {
  addTool(WindowLevelRegionTool);
  addTool(PlanarFreehandContourSegmentationTool);
  addTool(SegmentSelectTool);
+  addTool(LabelmapSlicePropagationTool);
+  addTool(MarkerLabelmapTool);
+  addTool(RegionSegmentPlusTool);
  // Modify annotation tools to use dashed lines on SR
  const annotationStyle = {
    textBoxFontSize: '15px',
@@ -152,6 +157,9 @@ const toolNames = {
  PlanarFreehandContourSegmentation: PlanarFreehandContourSegmentationTool.toolName,
  SegmentBidirectional: SegmentBidirectionalTool.toolName,
  SegmentSelect: SegmentSelectTool.toolName,
+  LabelmapSlicePropagation: LabelmapSlicePropagationTool.toolName,
+  MarkerLabelmap: MarkerLabelmapTool.toolName,
+  RegionSegmentPlus: RegionSegmentPlusTool.toolName,
 };

 export { toolNames };
@@ -360,7 +360,7 @@ class SegmentationService extends PubSubService {
        type: LABELMAP,
        data: {
          imageIds: segImageIds,
-          referencedVolumeId: this._getVolumeIdForDisplaySet(displaySet),
+          // referencedVolumeId: this._getVolumeIdForDisplaySet(displaySet),
          referencedImageIds: referenceImageIds,
        },
      },