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

extensions/cornerstone-dicom-pmap/package.json

@@ -46,8 +46,8 @@
   },
   "dependencies": {
     "@babel/runtime": "^7.20.13",
-    "@cornerstonejs/adapters": "^3.7.17",
-    "@cornerstonejs/core": "^3.7.17",
+    "@cornerstonejs/adapters": "^3.8.2",
+    "@cornerstonejs/core": "^3.8.2",
     "@kitware/vtk.js": "32.12.0",
     "react-color": "^2.19.3"
   }

extensions/cornerstone-dicom-seg/package.json

@@ -46,8 +46,8 @@
   },
   "dependencies": {
     "@babel/runtime": "^7.20.13",
-    "@cornerstonejs/adapters": "^3.7.17",
-    "@cornerstonejs/core": "^3.7.17",
+    "@cornerstonejs/adapters": "^3.8.2",
+    "@cornerstonejs/core": "^3.8.2",
     "@kitware/vtk.js": "32.12.0",
     "react-color": "^2.19.3"
   }

extensions/cornerstone-dicom-seg/src/getToolbarModule.ts

@@ -45,7 +45,7 @@ export function getToolbarModule({ servicesManager }: withAppTypes) {
         if (!toolNames) {
           return {
             disabled: false,
-            // isActive: true,
+            // isActive: false,
           };
         }
 

extensions/cornerstone-dicom-sr/package.json

@@ -46,9 +46,9 @@
   },
   "dependencies": {
     "@babel/runtime": "^7.20.13",
-    "@cornerstonejs/adapters": "^3.7.17",
-    "@cornerstonejs/core": "^3.7.17",
-    "@cornerstonejs/tools": "^3.7.17",
+    "@cornerstonejs/adapters": "^3.8.2",
+    "@cornerstonejs/core": "^3.8.2",
+    "@cornerstonejs/tools": "^3.8.2",
     "classnames": "^2.3.2"
   }
 }

extensions/cornerstone-dynamic-volume/package.json

@@ -42,8 +42,8 @@
   },
   "dependencies": {
     "@babel/runtime": "^7.20.13",
-    "@cornerstonejs/core": "^3.7.17",
-    "@cornerstonejs/tools": "^3.7.17",
+    "@cornerstonejs/core": "^3.8.2",
+    "@cornerstonejs/tools": "^3.8.2",
     "classnames": "^2.3.2"
   }
 }

extensions/cornerstone/package.json

@@ -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",

extensions/cornerstone/src/commandsModule.ts

@@ -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 {

extensions/cornerstone/src/customizations/segmentationPanelCustomization.tsx

@@ -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>
       );

extensions/cornerstone/src/initCornerstoneTools.js

@@ -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 };

extensions/cornerstone/src/services/SegmentationService/SegmentationService.ts

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

extensions/default/src/Toolbar/ToolBoxWrapper.tsx

@@ -42,14 +42,14 @@ export function ToolBoxButtonGroupWrapper({ groupId, buttonSection, ...props })
   );
 }
 
-export function ToolBoxButtonWrapper({ onInteraction, options, ...props }) {
+export function ToolBoxButtonWrapper({ onInteraction, className, options, ...props }) {
   return (
     <div className="bg-popover flex flex-row rounded-md px-0 py-0">
       <ToolButton
         {...props}
         id={props.id}
         size="small"
-        className={classNames(props.disabled && 'text-foreground/70')}
+        className={classNames(props.disabled && 'text-foreground/70', className)}
         onInteraction={event => {
           onInteraction?.({
             event,

extensions/default/src/ViewerLayout/index.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useState, useCallback } from 'react';
+import React, { useEffect, useState, useCallback, useRef } from 'react';
 import PropTypes from 'prop-types';
 
 import { InvestigationalUseDialog } from '@ohif/ui';
@@ -55,6 +55,18 @@ function ViewerLayout({
     setRightPanelClosed
   );
 
+  const leftPanelRef = useRef<HTMLDivElement>(null);
+  const rightPanelRef = useRef<HTMLDivElement>(null);
+
+  const handleMouseEnter = () => {
+    if (
+      leftPanelRef.current?.contains(document.activeElement) ||
+      rightPanelRef.current?.contains(document.activeElement)
+    ) {
+      (document.activeElement as HTMLElement)?.blur();
+    }
+  };
+
   const LoadingIndicatorProgress = customizationService.getCustomization(
     'ui.loadingIndicatorProgress'
   );
@@ -150,16 +162,17 @@ function ViewerLayout({
           {showLoadingIndicator && <LoadingIndicatorProgress className="h-full w-full bg-black" />}
           <ResizablePanelGroup {...resizablePanelGroupProps}>
             {/* LEFT SIDEPANELS */}
-
             {hasLeftPanels ? (
               <>
                 <ResizablePanel {...resizableLeftPanelProps}>
-                  <SidePanelWithServices
-                    side="left"
-                    isExpanded={!leftPanelClosedState}
-                    servicesManager={servicesManager}
-                    {...leftPanelProps}
-                  />
+                  <div ref={leftPanelRef}>
+                    <SidePanelWithServices
+                      side="left"
+                      isExpanded={!leftPanelClosedState}
+                      servicesManager={servicesManager}
+                      {...leftPanelProps}
+                    />
+                  </div>
                 </ResizablePanel>
                 <ResizableHandle
                   onDragging={onHandleDragging}
@@ -171,7 +184,10 @@ function ViewerLayout({
             {/* TOOLBAR + GRID */}
             <ResizablePanel {...resizableViewportGridPanelProps}>
               <div className="flex h-full flex-1 flex-col">
-                <div className="relative flex h-full flex-1 items-center justify-center overflow-hidden bg-black">
+                <div
+                  className="relative flex h-full flex-1 items-center justify-center overflow-hidden bg-black"
+                  onMouseEnter={handleMouseEnter}
+                >
                   <ViewportGridComp
                     servicesManager={servicesManager}
                     viewportComponents={viewportComponents}
@@ -188,12 +204,14 @@ function ViewerLayout({
                   className={resizableHandleClassName}
                 />
                 <ResizablePanel {...resizableRightPanelProps}>
-                  <SidePanelWithServices
-                    side="right"
-                    isExpanded={!rightPanelClosedState}
-                    servicesManager={servicesManager}
-                    {...rightPanelProps}
-                  />
+                  <div ref={rightPanelRef}>
+                    <SidePanelWithServices
+                      side="right"
+                      isExpanded={!rightPanelClosedState}
+                      servicesManager={servicesManager}
+                      {...rightPanelProps}
+                    />
+                  </div>
                 </ResizablePanel>
               </>
             ) : null}

extensions/default/src/utils/Toolbox.tsx

@@ -95,7 +95,7 @@ export function Toolbox({ buttonSectionId, title }: { buttonSectionId: string; t
         )}
       </PanelSection.Header>
 
-      <PanelSection.Content className="flex-shrink-0 border-none">
+      <PanelSection.Content className="bg-muted flex-shrink-0 border-none">
         {showConfig && <CustomConfigComponent />}
         {toolboxSections.map(section => {
           const sectionId = section.componentProps.buttonSection;
@@ -131,7 +131,7 @@ export function Toolbox({ buttonSectionId, title }: { buttonSectionId: string; t
           );
         })}
         {activeToolOptions && (
-          <div className="bg-primary-dark h-auto px-2">
+          <div className="bg-primary-dark mt-1 h-auto px-2">
             <ToolSettings options={activeToolOptions} />
           </div>
         )}

extensions/measurement-tracking/package.json

@@ -32,8 +32,8 @@
     "start": "yarn run dev"
   },
   "peerDependencies": {
-    "@cornerstonejs/core": "^3.7.17",
-    "@cornerstonejs/tools": "^3.7.17",
+    "@cornerstonejs/core": "^3.8.2",
+    "@cornerstonejs/tools": "^3.8.2",
     "@ohif/core": "3.10.0-beta.143",
     "@ohif/extension-cornerstone-dicom-sr": "3.10.0-beta.143",
     "@ohif/extension-default": "3.10.0-beta.143",